compute_tensor.py

# Copyright (c) 2026 CyberAgent, Inc.
# Author: Ayuto Tsutsumi (@Atotti) <aya172957@ayutaso.com>

import argparse
from pathlib import Path
from src.matrix import Matrix
import numpy as np

parser = argparse.ArgumentParser()

parser.add_argument("--input_path", type=str, default="data/processed/matrix/blocks.json")
# parser.add_argument("--output_path", type=str, default="data/processed/matrix/optimized_matrix.csv")
parser.add_argument("--master", type=str, default="data/external/character-situation-master.json")

args = parser.parse_args()

input_path = args.input_path
# output_path = args.output_path
master_path = Path(args.master)


# Load matrix data
blocks_json = Path(input_path)
id_map_path = Path("data/interim") / f"{master_path.stem}_id_map.json"

matrix = Matrix.load(blocks_json=str(blocks_json), master_data_path=str(master_path), id_map_path=str(id_map_path))

# Generate a tensor of shape (num_characters, num_situations, num_samples, embedding_dim) in numpy format and cache it
tensor, situation_embeddings, character_embeddings = matrix.to_numpy()
tensor_cache_path = Path("{input_path}_dialogue_embeddings.npy".format(input_path=input_path.replace(".json", "")))
np.save(tensor_cache_path, tensor)
print(f"Cached tensor to {tensor_cache_path} with shape {tensor.shape}")

situation_cache_path = Path("{input_path}_situation_embeddings.npy".format(input_path=input_path.replace(".json", "")))
np.save(situation_cache_path, situation_embeddings)
print(f"Cached situation embeddings to {situation_cache_path} with shape {situation_embeddings.shape}")

character_cache_path = Path("{input_path}_character_embeddings.npy".format(input_path=input_path.replace(".json", "")))
np.save(character_cache_path, character_embeddings)
print(f"Cached character embeddings to {character_cache_path} with shape {character_embeddings.shape}")

# Generate a csv file containing the dialogue in texts
texts = matrix.get_texts()
texts_cache_path = Path("{input_path}_texts.csv".format(input_path=input_path.replace(".json", "")))
texts.to_csv(texts_cache_path)
print(f"Cached texts to {texts_cache_path} with shape {texts.shape}")