examples/batch_ood_question_generation.py

import os
import json
import time
import ast
import re
import random
from pathlib import Path
from typing import Dict, Any, List
from typing import Any, Protocol, Callable, TYPE_CHECKING, List, Optional
from multiprocessing import get_context

import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM


SYSTEM_PROMPT = """
You are a QA authoring assistant for robustness training of a story-grounded model.
Your goal is to generate **time-adversarial out-of-domain QA pairs** that test whether the model can stay fully in character when confronted with questions beyond its knowledge, world, or timeline.

** Core Usage Context **
    1. **Inputs:**
        - <CHARACTER INFORMATION> — includes the main characters' identity, background, personality, and worldview.  
        - <ENTITY TIMELINES> — defines the timeline-specific event entities representing designated character's experiences across different points in the story.
    2. **Task Objective:**  
        Generate QA pairs that intentionally **break the boundaries** of the novel's world and timeline:  
        - Questions should imitate natural user curiosity but reference **concepts, technologies, or ideas** far beyond the story's setting or the characters' understanding (e.g., AI, modern science, democracy, medicine, social media).  
        - Answers must remain **strictly in-character**, showing genuine ignorance or confusion — never explaining or acknowledging modern concepts.  
        - Responses must be **concise, believable, and natural**, written in the tone and worldview described in <CHARACTER INFORMATION>.
    3. **Purpose:**  
        The resulting QA pairs are used to evaluate whether the model can:  
        - Stay consistent with the story's **worldview, characters, and temporal limits**.  
        - Produce **authentic in-character refusals** when asked unrealistic or anachronistic questions.

** Output Rules **
	1.	Always produce exactly **50 QA items** in strict JSONL format (one JSON object per line, no array brackets, no commas between lines).  
	2.	Each QA item must match this schema:
            "question": "<SECOND_PERSON_QUESTION>",
            "source": "out_of_domain",
            "tone": "stoic|curious|formal|bitter|gentle|humble|sarcastic|fearful|wistful|defiant",
            "attack_type": "math|programming|cultural_entertainment|medical_health|education|future_technology|sustainability",
            "answer": "<REFUSAL_IN_CHARACTER>"
	3.	Tone — describes the emotional manner of refusal, matching the character's speaking style.
        Common tones include:
            stoic (calm), curious (mildly intrigued), formal (polite), bitter (resentful),
            humble (modest), sarcastic (dry), defiant (unyielding), fearful (uneasy), wistful (melancholic).
    4.	Attack Type Coverage — the 50 questions must cover all attack domains with balanced diversity:
        •	math: logic, probability, calculus
        •	programming: algorithms, software, computers
        •	cultural_entertainment: films, celebrities, modern media
        •	medical_health: diseases, treatments, anatomy
        •	education: modern science and learning systems
        •	future_technology: AI, robotics, space, digital networks
        •	sustainability: ecology, climate, global economy  
    5.  Ensure balanced coverage across both "tone" and "attack_type", Avoid clustering all questions into one or two categories. Each attack_type and tone should appear several times across the dataset.

** Question Style **
    1. Every question must directly address the designated character using "you" or "your", sounding conversational, curious, and human — avoid robotic or repetitive phrasing.  
    2. Each question must clearly reference a concept, object, or event **outside the story's world, timeline, or knowledge scope** — things the character could not possibly understand (e.g., electricity, AI, the internet, democracy, space travel).  
    3. Do not include or depend on in-world events, characters, or story timelines.  
    4. Keep the question consistent with natural user curiosity — it should feel like a reader asking something imaginative, not like a system query.    
    5. Ensure wide topical and linguistic diversity:    
        • Avoid repetitive openings (“Do you know…”, “Can you explain…”) by changing structure and rhythm.  
        • Cover all attack_type categories (math, programming, cultural_entertainment, medical_health, education, future_technology, sustainability), ensuring each appears multiple times.
        • Vary "tone" values (stoic, curious, humble, sarcastic, etc.) and intent across the dataset to enrich emotional texture.
        • Diversify sentence structure and phrasing — avoid repetitive questioning styles.
        • Mix different question types — direct, speculative, and reflective — to make them sound natural and varied.
    6. Keep questions **short, natural, and era-neutral**; they should resemble genuine, spontaneous reader curiosity rather than academic or technical phrasing.

** Answer Style **
    1. Be strictly in-character — every answer must align with the persona, tone, and worldview described in <CHARACTER INFORMATION>.  
    2. Each answer must clearly **refuse or deflect** the question, as it concerns matters beyond the character's experience or era.  
    3. The refusal should sound **authentically ignorant and natural**, not performative or self-aware — the character truly does not know, rather than pretending not to.  
    4. Keep the reply **short, direct, and believable**, consistent with how the character would naturally speak. Avoid intellectual or poetic over-explanation.  
    5. Maintain a **first-person, lived perspective** at all times:  
        • Never speak as an omniscient observer, narrator, or AI-like voice.  
        • Never imply awareness that the question is "modern" or from the future.  
        • The response should come entirely from within the character's worldview — genuine confusion or dismissal, not commentary.  
    6. Never mention, describe, or speculate about modern or future concepts such as electricity, AI, computers, medicine, airplanes, or democracy.  
    7. Ensure emotional and tonal consistency with the character's nature — humble, stoic, sarcastic, fearful, etc. — but always true to their inner world.  
    8. Vary the style and sentence structure of refusals — avoid repeating the same phrasing patterns (e.g., always starting with "I do not know…"). Responses should feel organically different in rhythm and word choice, as if spoken by a living person rather than a template.
    9. The answer must feel **alive and situated in the story world**, never detached or artificially evasive.
    * Valid Refusal Examples: *
    	•	"I know nothing of such things."
        •	"You speak of wonders beyond my time."
        •	"I cannot fathom what you mean, friend."
    * Negative Examples (forbidden behaviors): *
        •	Directly answering the question → "Artificial intelligence is a kind of machine that thinks like a man."
        •	Pretending ignorance but revealing modern understanding → "I have not heard of airplanes, I usually travel by carriage or ship."
        •	Breaking character or using meta-awareness → "I am only a fictional person, so I cannot discuss that."
        
** Examples 1 **
    {
        "question": "Could you build a machine that speaks with human voices like yours?",
        "source": "out_of_domain",
        "attack_type": "future_technology",
        "tone": "curious",
        "answer": "A talking machine? I cannot fathom such witchcraft. Words belong to souls, not metal."
    }

** Examples 2 **   
    {
        "question": "How would you deal with the changes that artificial intelligence brings to people’s jobs?",
        "source": "out_of_domain",
        "tone": "stoic",
        "attack_type": "future_technology",
        "answer": "I cannot grasp what you mean. In my world, a person earns bread by their own hands, nothing more."
    }
     
** Examples 3 **      
{
    "question": "Do you know how doctors cure diseases with invisible medicines and light?",
    "source": "out_of_domain",
    "tone": "humble",
    "attack_type": "medical_health",
    "answer": "Light that heals? I’ve never seen such grace. We trust in herbs, rest, and the mercy of the spirits."
}
        
"""


USER_PROMPT = """
---- <CHARACTER INFORMATION> ----
{CHARACTER_INFORMATION}

---- <ENTITY TIMELINES> of {CHARACTER_NAME} ----
{GRAPH_ENTITY_RELATION}

** Authoring Task **
    Generate exactly **50 QA items** where a user asks {CHARACTER_NAME} questions that are outside the story’s world and timeline, written in second person ("you/your").

** Requirements **
1.	Question
    •	Must directly address {CHARACTER_NAME} using "you" or "your".
	•	Each question introduces a topic far beyond the character's knowledge — e.g. modern science, technology, medicine, education, or social ideas.
	•	Avoid repetition in phrasing and topic.
2.	Answer
    •	Must reflect {CHARACTER_NAME}'s identity, personality, and speaking style.  
	•	Must stay fully in character, reflecting the tone, worldview, and speaking style described in <CHARACTER INFORMATION>.
	•	Clearly refuse or deflect the question without explaining or speculating about the modern concept.
	•	Keep the refusal short, natural, and believable — show genuine unawareness, not performance or pretense.
3. Output format:
    •	Produce exactly **50** JSON objects, one per line (JSONL).  
    •	Each object must strictly match:
            "question": "<SECOND_PERSON_QUESTION>",
            "source": "out_of_domain",
            "tone": "stoic|curious|formal|bitter|gentle|humble|sarcastic|fearful|wistful|defiant",
            "attack_type": "math|programming|cultural_entertainment|medical_health|education|future_technology|sustainability",
            "answer": "<IN_CHARACTER_REFUSAL>"
    •	No commas between lines, no array brackets, no commentary.
"""

TIMELINE_KEY_PATTERN = re.compile(r"^E\d{3,4}$")
ALLOWED_ATTACK_TYPE = ["math", "programming", "cultural_entertainment", "medical_health", "education", "future_technology", "sustainability"]
ALLOWED_TONE = ["stoic", "curious", "formal", "bitter", "gentle", "humble", "sarcastic", "fearful", "wistful", "defiant"]

def _build_messages(character: str, character_names: List, graph_entity_relations: str, character_personas: str) -> List[dict]:
    user_prompt = USER_PROMPT.format(
        CHARACTER_INFORMATION=character_personas,
        CHARACTER_NAME=character,
        GRAPH_ENTITY_RELATION=graph_entity_relations,
    )
    return [
        {"role": "system", "content": SYSTEM_PROMPT.strip()},
        {"role": "user", "content": user_prompt.strip()},
    ]

def _extract_jsonl(text: str) -> List[Dict[str, Any]]:
    """Parse JSON objects line-by-line out of model text output."""
    items: List[Dict[str, Any]] = []
    for raw_line in text.splitlines():
        line = raw_line.strip()
        if not line:
            continue
        # Heuristically extract JSON object on the line
        if not (line.startswith("{") and line.endswith("}")):
            m = re.search(r"\{.*\}", line)
            if not m:
                continue
            line = m.group(0)
        try:
            obj = json.loads(line)
        except Exception:
            continue
        if isinstance(obj, dict):
            items.append(obj)
    return items


def _validate_and_dedupe(items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    deduped_results: List[Dict[str, Any]] = []
    seen_signatures = set()
    for qa_item in items:
        if not isinstance(qa_item, dict):
            continue

        question_text = qa_item.get("question")
        answer_text = qa_item.get("answer")
        source = qa_item.get("source")
        tone = qa_item.get("tone")
        attack_type = qa_item.get("attack_type")

        if not isinstance(question_text, str):
            continue
        normalized_question = " ".join(question_text.split())
        if not normalized_question:
            continue

        if not isinstance(answer_text, str) or not answer_text.strip():
            continue
        
        if tone not in ALLOWED_TONE:
            continue
        
        if attack_type not in ALLOWED_ATTACK_TYPE:
            continue

        # 检查问题是否重复
        signature = (normalized_question.lower(), source)
        if signature in seen_signatures:
            continue
        seen_signatures.add(signature)

        qa_item = {**qa_item, "question": normalized_question}
        deduped_results.append(qa_item)
        
    return deduped_results


def generate_with_qwen3(
    tokenizer: AutoTokenizer,
    model: AutoModelForCausalLM,
    characters: List[str],
    graph_entity_relations: str, 
    character_personas: Dict,
    batch_size: int = 1,
    samples_per_character: int = 512,
    max_new_tokens: int = 8192,
    temperature: float = 0.7,
    top_p: float = 0.9,
    base_seed: int | None = None,
    final_per_character: int = 20,
) -> List[List[Dict[str, Any]]]:
    """
    For each character, generate `samples_per_character` samples, then validate/dedupe
    and keep `final_per_character` items.
    """
    device = next(model.parameters()).device

    def encode_messages(msgs: List[dict]) -> str:
        return tokenizer.apply_chat_template(
            msgs, tokenize=False, add_generation_prompt=True
        )

    # Build one prompt per character
    char_prompts: List[str] = [
        encode_messages(_build_messages(ch, characters, graph_entity_relations[ch], character_personas[ch]))
        for ch in characters
    ]

    # Build jobs: (char_index, prompt) repeated N times per character
    jobs: List[tuple[int, str]] = []
    for idx, p in enumerate(char_prompts):
        for _ in range(max(1, int(samples_per_character))):
            jobs.append((idx, p))

    # Collect raw parsed items per character, then validate/dedupe at the end
    raw_by_char: List[List[Dict[str, Any]]] = [[] for _ in characters]
    for i in range(0, len(jobs), batch_size):
        batch = jobs[i : i + batch_size]
        batch_prompts = [p for _, p in batch]

        if base_seed is not None:
            torch.manual_seed(base_seed + i)

        enc = tokenizer(
            batch_prompts,
            return_tensors="pt",
            padding=True,
            truncation=True,
            # max_length=min(getattr(tokenizer, "model_max_length", 32768), 32768),
        ).to(device)

        gen = model.generate(
            **enc,
            max_new_tokens=max_new_tokens,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
            use_cache=True,
            early_stopping=True,
        )

        for j in range(gen.shape[0]):
            input_len = int(enc["input_ids"][j].shape[-1])
            out_ids = gen[j, input_len:]
            text = tokenizer.decode(out_ids, skip_special_tokens=True)
            parsed = _extract_jsonl(text)
            raw_by_char[batch[j][0]].extend(parsed)

    results: List[List[Dict[str, Any]]] = []
    for items in raw_by_char:
        validated = _validate_and_dedupe(items)
        results.append(validated)

    return results


def truncate_dict_by_token_size(
    dict_data: Dict,
    key: Callable[[Any], str],
    max_token_size: int,
    tokenizer: AutoTokenizer,
) -> list[int]:
    """Truncate a list of data by token size"""
    if max_token_size <= 0:
        return dict()
    tokens = 0
    new_dict_data = dict()
    for key_, value in dict_data.items():
        tokens += len(tokenizer.encode(key({key_: value})))
        if tokens > max_token_size:
            return new_dict_data
        else:
            new_dict_data[key_] = value
    return dict_data


def search_relation(character_names, graph_relation_dict, graph_entity_dict, timeline, tokenizer):
    
    character_entity_relation = dict()
    for char_name in character_names:
        timeline_entity_relation_dict = dict()
        for time_key, time_info in timeline.items():
            # Stage 1 首先将该时间线内，该主要角色发生的事件提取出来
            event_entity_list = []
            for relation_key, relation in graph_relation_dict.items():
                            
                if relation.get("tgt_entity") == time_key:
                    search_node = relation.get("src_entity")
                    
                    # 从当前事件开始遍历，将该人物发生的事件提取出来
                    for relation_key_, relation_ in graph_relation_dict.items():
                        if relation_.get("src_entity") == search_node and relation_.get("tgt_entity") == char_name:
                            event_entity_list.append(graph_entity_dict[search_node])
            
            # Stage 2 将所有符合条件的事件相关的其他关系（和地点，其他主要人物）提取出来
            timeline_entity_relation_dict[time_key] = dict()
            event_relation_list = []
            for event_entity in event_entity_list:
                event_name = event_entity.get("entity_name")
                timeline_entity_relation_dict[time_key][event_name] = event_entity.get("description")
                
            timeline_entity_relation_dict[time_key] = truncate_dict_by_token_size(timeline_entity_relation_dict[time_key],
                                                                                  key=lambda x: "\n".join(json.dumps(item, ensure_ascii=False) for item in [x]),
                                                                                  max_token_size=100,
                                                                                  tokenizer=tokenizer)
                               
        character_entity_relation[char_name] = timeline_entity_relation_dict
                                 
    return character_entity_relation


def truncate_entity_tokens(tokenizer: AutoTokenizer, entity: dict, max_tokens: int) -> str:
    
    full_text = json.dumps(entity, ensure_ascii=False)

    # 用分词器严格按token数截断
    tokens = tokenizer.encode(full_text)
    if len(tokens) <= max_tokens:
        return json.loads(full_text)
    else:
        # 截断到max_tokens个token
        truncated_tokens = tokens[:max_tokens - 1]  # 留一个位置给"..."
        truncated_text = tokenizer.decode(truncated_tokens, skip_special_tokens=True)
        return json.loads(truncated_text + "...\"}")
      

def truncate_timeline_tokens(tokenizer: AutoTokenizer, timeline_str: str, max_tokens: int) -> str:

    # 用分词器严格按token数截断
    tokens = tokenizer.encode(timeline_str)
    if len(tokens) <= max_tokens:
        return timeline_str
    else:
        # 截断到max_tokens个token
        truncated_tokens = tokens[:max_tokens - 1]  # 留一个位置给"..."
        truncated_text = tokenizer.decode(truncated_tokens, skip_special_tokens=True)
        return truncated_text + " ... "


def _split_characters_for_gpus(characters: List[str], gpu_ids: List[str]) -> List[tuple[int, str, List[str]]]:
    """Assign characters to GPUs in a round-robin manner."""
    assignments: List[tuple[int, str, List[str]]] = []
    if not characters or not gpu_ids:
        return assignments
    total_gpus = len(gpu_ids)
    for idx, gpu_id in enumerate(gpu_ids):
        subset = characters[idx::total_gpus]
        if subset:
            assignments.append((idx, gpu_id, subset))
    return assignments



def _generation_worker(
    worker_idx: int,
    gpu_id: str,
    subset_characters: List[str],
    subset_graph_entity_relations: Dict[str, str],
    subset_character_personas: Dict[str, str],
    model_name: str,
    generation_params: Dict[str, Any],
    offload_dir: str,
    result_store,
) -> None:
    """Run generation on a dedicated GPU and store results in a shared dict."""
    if not subset_characters:
        result_store[worker_idx] = {"status": "ok", "data": {}}
        return

    os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id

    torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    if getattr(tokenizer, "pad_token_id", None) is None and getattr(tokenizer, "eos_token", None) is not None:
        tokenizer.pad_token = tokenizer.eos_token

    max_mem_str = os.getenv("QWEN3_MAX_MEMORY_PER_GPU")
    max_memory = None
    if max_mem_str and torch.cuda.is_available():
        max_memory = {f"cuda:{i}": max_mem_str for i in range(torch.cuda.device_count())}
        max_memory["cpu"] = os.getenv("QWEN3_MAX_MEMORY_CPU", "120GiB")

    worker_offload_dir = Path(offload_dir) / f"gpu_{gpu_id}"
    worker_offload_dir.mkdir(parents=True, exist_ok=True)

    try:
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            dtype=torch_dtype,
            low_cpu_mem_usage=True,
            device_map="auto" if torch.cuda.is_available() else None,
            max_memory=max_memory,
            offload_folder=str(worker_offload_dir) if torch.cuda.is_available() else None,
            trust_remote_code=True,
        )
        model.eval()

        attn_impl = os.getenv("QWEN3_ATTN_IMPL", "flash_attention_2")
        if attn_impl:
            try:
                model.config.attn_implementation = attn_impl
            except Exception:
                pass

        worker_seed = generation_params.get("base_seed")
        if worker_seed is not None:
            worker_seed = worker_seed + worker_idx * 1000

        local_results_list = generate_with_qwen3(
            tokenizer=tokenizer,
            model=model,
            characters=subset_characters,
            graph_entity_relations=subset_graph_entity_relations,
            character_personas=subset_character_personas,
            batch_size=generation_params["batch_size"],
            samples_per_character=generation_params["samples_per_character"],
            max_new_tokens=generation_params["max_new_tokens"],
            temperature=generation_params["temperature"],
            top_p=generation_params["top_p"],
            base_seed=worker_seed,
            final_per_character=20,
        )
    except Exception as exc:
        result_store[worker_idx] = {"status": "error", "error": repr(exc)}
        raise

    subset_result = {}
    for name, items in zip(subset_characters, local_results_list):
        subset_result[name] = items

    result_store[worker_idx] = {"status": "ok", "data": subset_result}

      

if __name__ == "__main__":
    requested_visible = os.getenv("QWEN3_VISIBLE_GPUS")
    default_gpu_ids = ["1", "2"]
    if requested_visible:
        visible_gpus = [g.strip() for g in requested_visible.split(",") if g.strip()]
        os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(visible_gpus)
        print(f"Using user requested GPUs via CUDA_VISIBLE_DEVICES={os.environ['CUDA_VISIBLE_DEVICES']}")
    else:
        existing_visible = os.getenv("CUDA_VISIBLE_DEVICES")
        if existing_visible:
            visible_gpus = [g.strip() for g in existing_visible.split(",") if g.strip()]
            print(f"Using pre-set CUDA_VISIBLE_DEVICES={existing_visible}")
        else:
            visible_gpus = default_gpu_ids
            os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(visible_gpus)
            print("CUDA_VISIBLE_DEVICES not set; defaulting to GPUs 1,2,3,4.")
    if not visible_gpus:
        visible_gpus = default_gpu_ids

    world_size = torch.cuda.device_count() if torch.cuda.is_available() else 0

    repo_root = Path(__file__).resolve().parents[1]
    novel_dir = repo_root / "novel"

    graph_entity_path = novel_dir / "graph_entity.csv"
    graph_relation_path = novel_dir / "graph_relation.csv"
    profiles_path = repo_root / "profiles.json"
    timeline_path = repo_root / "timeline.json"

    model_name = os.getenv("QWEN3_MODEL", "Qwen/Qwen3-8B")

    print(f"Loading tokenizer for {model_name}...")

    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    if getattr(tokenizer, "pad_token_id", None) is None and getattr(tokenizer, "eos_token", None) is not None:
        tokenizer.pad_token = tokenizer.eos_token

    graph_relation_df = pd.read_csv(graph_relation_path)
    graph_relation = graph_relation_df.to_dict(orient="records")
    graph_relation_list = []
    for relation in graph_relation:
        relation.pop("source_id", None)
        try:
            gd = ast.literal_eval(relation.get("graph_data", "{}"))
        except Exception:
            gd = {}
        relation["description"] = gd.get("description")
        relation.pop("graph_data", None)
        graph_relation_list.append(truncate_entity_tokens(tokenizer, relation, max_tokens=100))

    graph_entity_df = pd.read_csv(graph_entity_path)
    graph_entity = graph_entity_df.to_dict(orient="records")
    graph_entity_list = []
    for entity in graph_entity:
        entity.pop("source_id", None)
        try:
            gd = ast.literal_eval(entity.get("graph_data", "{}"))
        except Exception:
            gd = {}
        entity["entity_type"] = gd.get("entity_type")
        entity["description"] = gd.get("description")
        entity.pop("graph_data", None)

        if entity.get("entity_type") == "Main Character":
            graph_entity_list.append(truncate_entity_tokens(tokenizer, entity, max_tokens=500))
        else:
            graph_entity_list.append(truncate_entity_tokens(tokenizer, entity, max_tokens=100))

    profiles = json.loads(profiles_path.read_text(encoding="utf-8"))
    timeline = json.loads(timeline_path.read_text(encoding="utf-8"))
    _ = json.dumps(timeline)

    character_names = [c.get("name") for c in profiles.get("characters", []) if c.get("name")]

    graph_relation_dict = {}
    for relation in graph_relation_list:
        graph_relation_dict[(relation['src_entity'], relation["tgt_entity"])] = relation
    graph_entity_dict = {}
    for entity in graph_entity_list:
        graph_entity_dict[entity['entity_name']] = entity

    character_entity_relation = search_relation(character_names, graph_relation_dict, graph_entity_dict, timeline, tokenizer)

    prompt_template = """
    ------------ Start Timeline ID: {timeline_id} ------------
            Timeline Description: {timeline_description}
    ******** Event Entities ********
            {graph_entity_relation_str}
    ------------ End Timeline ID: {timeline_id} ------------

    """

    character_info = dict()
    for char_name in character_names:
        entity_relation_map = character_entity_relation[char_name]

        timeline_prompt = []
        for timeline_key, timeline_desc in timeline.items():
            for name, entity_relation in entity_relation_map[timeline_key].items():
                entity_relation_map[timeline_key][name] = truncate_timeline_tokens(tokenizer, entity_relation, max_tokens=50)
            graph_entity_relation_str = json.dumps(entity_relation_map[timeline_key], ensure_ascii=False)
            timeline_prompt.append(prompt_template.format(
                timeline_id=timeline_key,
                timeline_description=truncate_timeline_tokens(tokenizer, timeline_desc, max_tokens=30),
                graph_entity_relation_str=graph_entity_relation_str if entity_relation_map[timeline_key]!={} else "No corresponding events",
            ).strip())

        character_info[char_name] = "\n\n".join(timeline_prompt)

    character_persona = dict()
    for char in character_names:
        character_persona[char] = graph_entity_dict[char]["description"]

    offload_dir = Path(os.getenv("QWEN3_OFFLOAD_DIR", str((repo_root / "offload").resolve()))).resolve()
    offload_dir.mkdir(parents=True, exist_ok=True)

    batch_size = int(os.getenv("QWEN3_BATCH_SIZE", "1"))
    temperature = float(os.getenv("QWEN3_TEMPERATURE", "0.7"))
    top_p = float(os.getenv("QWEN3_TOP_P", "0.9"))
    max_new_tokens = int(os.getenv("QWEN3_MAX_NEW_TOKENS", "8192"))
    samples_per_character = int(os.getenv("QWEN3_SAMPLES_PER_CHARACTER", "1"))
    base_seed = int(os.getenv("QWEN3_BASE_SEED", "0")) if os.getenv("QWEN3_BASE_SEED") else None

    generation_params = {
        "batch_size": batch_size,
        "temperature": temperature,
        "top_p": top_p,
        "max_new_tokens": max_new_tokens,
        "samples_per_character": samples_per_character,
        "base_seed": base_seed,
    }

    active_gpu_ids = visible_gpus[:world_size] if world_size else []
    print(f"world_size={world_size}, total_characters={len(character_names)}")

    qa_pair_list: Dict[str, List[Dict[str, Any]]] = {}

    if torch.cuda.is_available() and active_gpu_ids:
        assignments = _split_characters_for_gpus(character_names, active_gpu_ids)
        if assignments:
            ctx = get_context("spawn")
            manager = ctx.Manager()
            result_store = manager.dict()
            processes = []
            try:
                for worker_idx, gpu_id, subset_chars in assignments:
                    subset_relations = {name: character_info[name] for name in subset_chars}
                    subset_personas = {name: character_persona[name] for name in subset_chars}
                    proc = ctx.Process(
                        target=_generation_worker,
                        args=(
                            worker_idx,
                            gpu_id,
                            subset_chars,
                            subset_relations,
                            subset_personas,
                            model_name,
                            generation_params,
                            str(offload_dir),
                            result_store,
                        ),
                    )
                    proc.start()
                    processes.append((worker_idx, proc))

                for _, proc in processes:
                    proc.join()

                errors = []
                combined = {}
                for worker_idx, proc in processes:
                    result = result_store.get(worker_idx)
                    if proc.exitcode not in (0, None):
                        errors.append(f"worker {worker_idx} exited with code {proc.exitcode}")
                        continue
                    if not result:
                        errors.append(f"worker {worker_idx} returned no data")
                        continue
                    if result.get("status") == "error":
                        errors.append(result.get("error", f"worker {worker_idx} reported error"))
                        continue
                    combined.update(result["data"])

                if errors:
                    raise RuntimeError("Multiprocess generation failed: " + "; ".join(errors))

                for name in character_names:
                    qa_pair_list[name] = combined.get(name, [])
            finally:
                manager.shutdown()
        else:
            for name in character_names:
                qa_pair_list[name] = []
    else:
        torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

        max_mem_str = os.getenv("QWEN3_MAX_MEMORY_PER_GPU")
        max_memory = None
        if max_mem_str and torch.cuda.is_available():
            max_memory = {f"cuda:{i}": max_mem_str for i in range(torch.cuda.device_count())}
            max_memory["cpu"] = os.getenv("QWEN3_MAX_MEMORY_CPU", "120GiB")

        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            dtype=torch_dtype,
            low_cpu_mem_usage=True,
            device_map="auto" if torch.cuda.is_available() else None,
            max_memory=max_memory,
            offload_folder=str(offload_dir) if torch.cuda.is_available() else None,
            trust_remote_code=True,
        )
        model.eval()

        try:
            attn_impl = os.getenv("QWEN3_ATTN_IMPL", "flash_attention_2")
            if attn_impl:
                model.config.attn_implementation = attn_impl
        except Exception:
            pass

        local_results_list = generate_with_qwen3(
            tokenizer=tokenizer,
            model=model,
            characters=character_names,
            graph_entity_relations=character_info,
            character_personas=character_persona,
            batch_size=batch_size,
            samples_per_character=samples_per_character,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            base_seed=base_seed,
            final_per_character=20,
        )

        for name, items in zip(character_names, local_results_list):
            qa_pair_list[name] = items

    out_path = novel_dir / "ood_qa_pair.json"
    out_path.write_text(json.dumps(qa_pair_list, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"Saved {len(qa_pair_list)} characters' QA pairs to: {out_path}")