TalkMatrix

TalkMatrix is a method for generating diverse and consistent character dialogues for interactive narrative applications (games, visual novels, etc.). Given a set of characters and situations, it:

1. Generates multiple candidate dialogues per (character, situation) pair using a local LLM.
2. Embeds each dialogue with both a style embedding and a semantic (text) embedding.
3. Optimises the selection of one dialogue per cell via JAX-accelerated coordinate ascent to maximise four objectives: style diversity, style consistency, content diversity, and content consistency.
4. Evaluates the resulting dialogue matrix using an LLM-as-a-judge.

All four steps run fully locally — no API credentials are required.

---

Table of Contents

• Installation
• Data Format
• Quick Start (Example Data)
• Step-by-Step Guide
  • 1. Prepare Master Data
  • 2. Generate Dialogue Samples
  • 3. Compute Embeddings
  • 4. Cache Tensors
  • 5. Optimise
  • 6. Evaluate
• Configuration Reference
• Objectives

---

Installation

Prerequisites

• Python ≥ 3.11
• uv (recommended) or pip
• A CUDA-capable GPU is strongly recommended for step 1 (generation) and step 5 (JAX optimisation). CPU-only execution is possible but slow for large models.

Install dependencies

uv sync --dev

Or with pip:

pip install -r requirements.txt

Note: JAX GPU support requires an additional installation step.
Run pip install jax[cuda12] (or the appropriate CUDA variant) to enable GPU acceleration for the optimiser.

---

Data Format

TalkMatrix uses a master JSON file that describes the characters and situations for a scenario. An example is provided at example_data/example.json.

{
  "character": [
    {
      "Name": "キャラクター名",
      "Intro": "背景・存在:\n...\n性格・特徴:\n...\n口調・一人称:\n...",
      "ref_dialogues": [
        "そのキャラクターらしいセリフ1",
        "そのキャラクターらしいセリフ2",
        "そのキャラクターらしいセリフ3"
      ]
    }
  ],
  "situation": [
    {
      "Name": "シチュエーション名",
      "Intro": "シチュエーションの説明文。"
    }
  ]
}

Field	Description
character[].Name	Unique character name
character[].Intro	Character background, personality, and speech-style description
character[].ref_dialogues	1–3 reference lines in the character's voice (used for style embedding)
situation[].Name	Unique situation name
situation[].Intro	Situation description (used for content/consistency scoring)

---

Quick Start (Example Data)

The repository ships with a small ready-to-use scenario at example_data/example.json (3 Japanese historical characters × 4 board-game situations). Run the full pipeline in one command:

bash run_local.sh

This will:

1. Generate 4 dialogue samples per cell using Qwen/Qwen3-1.7B (local).
2. Embed all dialogues with MStyleDistance (style) and sbintuitions/sarashina-embedding-v1-1b (semantic).
3. Embed character and situation master data with the same local models.
4. Cache all tensors to .npy files for fast optimisation.
5. Run JAX coordinate-ascent optimisation (minimax/minimax) and save the optimised matrix.
6. Evaluate the optimised vs. random matrix with Qwen/Qwen3-1.7B as the LLM judge.

Results are saved to data/processed/example/.

Custom parameters

bash run_local.sh \
    example_data/example.json \   # master JSON
    data/processed/my_run \       # output directory
    Qwen/Qwen3-1.7B \             # generation model (HuggingFace id)
    Qwen/Qwen3-1.7B \             # judge model (HuggingFace id)
    8                             # samples per cell

---

Step-by-Step Guide

1. Prepare Master Data

Place your master JSON at example_data/<stem>.json (or anywhere you prefer).
Run update_id_map to initialise the internal ID mapping:

# The ID map is created automatically during the generate step.
# No manual action needed for new scenarios.

2. Generate Dialogue Samples

python3 main.py generate \
    --num 8 \
    --output_dir data/processed/my_scenario \
    --model_name Qwen/Qwen3-1.7B \
    --temperature 0.9 \
    --master example_data/my_scenario.json

Argument	Default	Description
--num	1	Number of candidate dialogues per (character, situation) cell
--model_name	Qwen/Qwen3-1.7B	HuggingFace model id (with /) for local inference
--temperature	0.7	Sampling temperature
--output_dir	data/processed/matrix	Directory for the generated blocks.json and cell embeddings
--master	data/external/...	Path to the master JSON file

Local model generation

Model name format	Backend used
Contains / (e.g. Qwen/Qwen3-1.7B)	Local HuggingFace Transformers
calm in name	Local CALM3 (HuggingFace)

3. Compute Embeddings

Two separate embedding steps are needed:

3a. Style embeddings (always local)

Uses StyleDistance/mstyledistance (768-dim, multilingual).

python3 main.py style_embedding \
    --input_path data/processed/my_scenario/blocks.json \
    --output_dir data/processed/my_scenario \
    --master example_data/my_scenario.json

3b. Text/semantic embeddings

python3 main.py text_embedding \
    --input_path data/processed/my_scenario/blocks.json \
    --output_dir data/processed/my_scenario \
    --id_map_path data/interim/my_scenario_id_map.json \
    --master example_data/my_scenario.json

Text/semantic embeddings always use the local SentenceTransformerEmbedding model (sbintuitions/sarashina-embedding-v1-1b by default, or override with LOCAL_TEXT_EMBEDDING_MODEL).

3c. Master character/situation embeddings

python3 main.py master_embedding \
    --input_path data/processed/my_scenario/blocks.json \
    --output_dir data/processed/my_scenario \
    --master example_data/my_scenario.json

This computes:

• Character: style_ref_embedding (MStyleDistance on ref_dialogues) + background_embedding (text model on Intro)
• Situation: content_embedding (text model on Intro)

Important: steps 3b and 3c must use the same text embedding model (set via LOCAL_TEXT_EMBEDDING_MODEL).

4. Cache Tensors

Assembles all embeddings into contiguous .npy tensors for fast JAX loading:

python3 compute_tensor.py \
    --input_path data/processed/my_scenario/blocks.json \
    --master example_data/my_scenario.json

This writes three files:

• blocks_dialogue_embeddings.npy — shape (n_situations, n_characters, n_samples, emb_dim)
• blocks_situation_embeddings.npy — shape (n_situations, emb_dim)
• blocks_character_embeddings.npy — shape (n_characters, emb_dim)

5. Optimise

Run the JAX coordinate-ascent optimiser (recommended: minimax aggregation and combination):

python3 main.py optimize \
    --input_path data/processed/my_scenario/blocks.json \
    --output_path data/processed/my_scenario/optimized_matrix.csv \
    --mode optuna_sum \
    --weights 1.0 1.0 1.0 1.0 \
    --n_trials 50 \
    --seed 42 \
    --master example_data/my_scenario.json

Or use the standalone coordinate-ascent script for more control:

from src.optimization.jax import coordinate_ascent_optimizer
import numpy as np

tensor = np.load("data/processed/my_scenario/blocks_dialogue_embeddings.npy")
sit_emb = np.load("data/processed/my_scenario/blocks_situation_embeddings.npy")
char_emb = np.load("data/processed/my_scenario/blocks_character_embeddings.npy")

index_matrices, component_values = coordinate_ascent_optimizer(
    dialogue_data=tensor,
    situation_data=sit_emb,
    character_data=char_emb,
    weights=(1.0, 1.0, 1.0, 1.0),
    restarts=20,
    agg="minimax",
    combine="minimax",
    seed=42,
)

6. Evaluate

Local LLM judge (no credentials needed)

python3 main.py evaluate \
    --judge_model Qwen/Qwen3-1.7B

---

Configuration Reference

Environment Variables

Variable	Default	Description
LOCAL_TEXT_EMBEDDING_MODEL	sbintuitions/sarashina-embedding-v1-1b	HuggingFace model id for local text embeddings

Choosing a Local Text Embedding Model

Any model compatible with sentence-transformers can be used. The default (sbintuitions/sarashina-embedding-v1-1b) is optimised for Japanese text. For a smaller footprint:

export LOCAL_TEXT_EMBEDDING_MODEL=paraphrase-multilingual-MiniLM-L12-v2

Consistency requirement: use the same LOCAL_TEXT_EMBEDDING_MODEL value for all embedding steps (3b, 3c) and the tensor caching step (4).

---

Objectives

The optimiser maximises a weighted combination of four components:

Symbol	Name	Description
SD	Style Diversity	Style variation across characters within each situation
SC	Style Consistency	Style coherence of each character across situations + alignment with ref_dialogues
CD	Content Diversity	Semantic variation across characters within each situation
CC	Content Consistency	Semantic alignment of each dialogue with its situation description

The combine mode controls how components are merged:

• sum — weighted sum (default, most permissive)
• harmonic — harmonic mean (penalises weak components)
• minimax — maximise the worst component (most conservative, recommended for balanced quality)

The agg mode controls how per-cell scores are pooled:

• mean — arithmetic mean
• harmonic — harmonic mean
• minimax — minimum (safeguards the weakest pair)

LICENSE

MIT License

Authors

Ayuto Tsutsumi @Atotti (aya172957@ayutaso.com)
Yuu Jinnai @jinnaiyuu (jinnai_yu@cyberagent.co.jp)

Reference

TBA