TalkMatrix/optimize_matrix.py at master · CyberAgentAILab/TalkMatrix · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
# Copyright (c) 2026 CyberAgent AI Lab
# Author: Ayuto Tsutsumi (@Atotti) <aya172957@ayutaso.com>

import argparse
from pathlib import Path
import random

import numpy as np
import pandas as pd


from src.optimization.jax import optimizer

if __name__ == "__main__":
    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")
    parser.add_argument(
        "--mode",
        type=str,
        choices=[
            "powell",
            "random",
            "optuna",
            "optuna_sum",
            "optuna_local_global",
            "optuna_local_consistency_diversity",
            "optuna_local_style_content",
        ],
        default="optuna",
        help="Optimization method to use",
    )
    parser.add_argument(
        "--weights",
        type=float,
        nargs=4,
        default=[1.0, 1.0, 1.0, 1.0],
        help="Weights for Situation Diversity (SD), Situation Consistency (SC),\
    Content Diversity (CD), and Content Consistency (CC)",
        metavar=("SD", "SC", "CD", "CC"),
    )
    parser.add_argument("--n_trials", type=int, default=50, help="Number of trials for optuna optimization")
    parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility")

    args = parser.parse_args()

    input_path = args.input_path
    output_path = args.output_path
    mode = args.mode
    weights = tuple(args.weights)
    n_trials = args.n_trials
    seed = args.seed

    # Note that fixing random seeds does not guarantee complete reproducibility when running in parallel.
    # For debugging or exact reproducibility, run in a single process.
    random.seed(seed)
    np.random.seed(seed)  # Legacy code might use numpy's global random seed

    # load dialogue embeddings tensor, situation embeddings, and character embeddings
    tensor_cache_path = Path("{input_path}_dialogue_embeddings.npy".format(input_path=input_path.replace(".json", "")))
    tensor = np.load(tensor_cache_path)
    print(f"Loaded tensor from {tensor_cache_path} with shape {tensor.shape}")

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

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

    texts_cache_path = Path("{input_path}_texts.csv".format(input_path=input_path.replace(".json", "")))
    texts = pd.read_csv(texts_cache_path)

    obj_name = mode.replace("optuna_", "")

    index_matrices, component_values = optimizer(
        dialogue_data=tensor,
        situation_data=situation_embeddings,
        character_data=character_embeddings,
        weights=weights,
        n_trials=n_trials,
        obj_name=obj_name,
    )
    for i in range(len(index_matrices)):
        index_matrix = index_matrices[i]
        component_value = component_values[i]

        optimization_config = Path(output_path.replace(".csv", f"_{obj_name}_optimized_{i}.csv") + ".config")
        index_matrix_str = "\n".join([str(index) for index in index_matrix])

        # Format component values for config file
        component_value_str = "\n".join([f"{k}={v}" for k, v in component_value.items()])

        with optimization_config.open("w", encoding="utf-8") as f:
            f.write(f"{mode=}\n{weights=}\n\n{component_value_str}\n\nindex_matrix=\n{index_matrix_str}")

        index_matrix_np = np.array(index_matrix)
        np.save(output_path.replace(".csv", f"_{obj_name}_optimized_{i}.npy"), index_matrix_np)

        # Retrieve the texts corresponding to the optimized index matrix

        optimized_texts = []
        for sit_idx in range(index_matrix_np.shape[0]):
            for char_idx in range(index_matrix_np.shape[1]):
                sample_idx = index_matrix_np[sit_idx, char_idx]

                entry = texts[
                    (texts["character_index"] == char_idx)
                    & (texts["situation_index"] == sit_idx)
                    & (texts["dialogue_index"] == sample_idx)
                ]
                data = {
                    "character_name": entry["character"].values[0],
                    "situation_name": entry["situation"].values[0],
                    "dialogue": entry["dialogue"].values[0],
                }
                optimized_texts.append(data)
        optimized_texts_df = pd.DataFrame(optimized_texts)
        optimized_texts_df.to_csv(output_path.replace(".csv", f"_{obj_name}_optimized_{i}.csv"), index=False)
        print(f"Saved optimized texts to {output_path.replace('.csv', f'_{obj_name}_optimized_{i}.csv')}")