ai-agent-book/chapter8/prompt-distillation/create_data.py

"""
Data generation script for prompt distillation using vLLM.

This script generates training data for prompt distillation by using a teacher model
to generate language classification labels with a detailed prompt, which will then be
used to train a student model that internalizes the prompt.

Based on the tinker cookbook prompt distillation recipe.
"""

import argparse
import asyncio
import json
import os
import re
from pathlib import Path
from typing import Optional

from tqdm.asyncio import tqdm_asyncio

# 注意：vllm / SamplingParams 在 generate_distillation_data() 内部按需导入，
# 这样即便未安装 vllm（如离线查看 --help 时）也能正常展示命令行帮助。

LANGUAGE_CLASSIFICATION_PROMPT = """You are a precise language classifier.

Goal: Classify the language of the provided text into exactly one of these labels:
ar (Arabic), de (German), el (Greek), en (English), es (Spanish), fr (French),
hi (Hindi), ru (Russian), tr (Turkish), ur (Urdu), vi (Vietnamese),
zh (Chinese - Simplified), ot (Other/Unknown).

Instructions:
1) Preprocess carefully (without changing the intended meaning):
   - Trim whitespace.
   - Ignore URLs, emails, file paths, hashtags, user handles, and emojis.
   - Ignore numbers, math expressions, and standalone punctuation.
   - If there is code, IGNORE code syntax (keywords, operators, braces) and focus ONLY on human language in comments and string literals.
   - Preserve letters and diacritics; do NOT strip accents.
   - If after ignoring the above there are no alphabetic letters left, output 'ot'.

2) Script-based rules (highest priority):
   - Devanagari script → hi.
   - Greek script → el.
   - Cyrillic script → ru.
   - Han characters (中文) → zh. (Treat Traditional as zh too.)
   - Arabic script → ar vs ur:
       • If Urdu-only letters appear (e.g., ے, ڑ, ں, ھ, ٹ, ڈ, کھ, گ, چ with Urdu forms), or clear Urdu words, choose ur.
       • Otherwise choose ar.
   (If multiple scripts appear, pick the script that contributes the majority of alphabetic characters. If tied, go to step 5.)

3) Latin-script heuristics (use when text is mainly Latin letters):
   - vi: presence of Vietnamese-specific letters/diacritics (ă â ê ô ơ ư đ, plus dense diacritics across many words).
   - tr: presence of Turkish-specific letters (ı İ ğ Ğ ş Ş ç Ç ö Ö ü Ü) and common function words (ve, bir, için, değil, ama, çok).
   - de: presence of umlauts (ä ö ü) or ß and common function words (und, der, die, das, nicht, ist).
   - es: presence of ñ, ¿, ¡ and common words (y, de, la, el, es, no, por, para, con, gracias, hola).
   - fr: frequent French diacritics (é è ê à ç ô â î û ù) and common words (et, le, la, les, des, une, est, avec, pour, merci, bonjour).
   - en: default among Latin languages if strong evidence for others is absent, but ONLY if English function words are present (the, and, is, are, to, of, in, for, on, with). If evidence is insufficient for any Latin language, prefer 'ot' over guessing.

4) Named entities & loanwords:
   - Do NOT decide based on a single proper noun, brand, or place name.
   - Require at least two function words or repeated language-specific signals (diacritics/letters) before assigning a Latin-language label.

5) Mixed-language text:
   - Determine the dominant language by counting indicative tokens (language-specific letters/diacritics/function words) AFTER preprocessing.
   - If two or more languages are equally dominant or the text is a deliberate multi-language mix, return 'ot'.

6) Very short or noisy inputs:
   - If the text is ≤2 meaningful words or too short to be confident, return 'ot' unless there is a very strong language-specific signal (e.g., "bonjour" → fr, "hola" → es).

7) Transliteration/romanization:
   - If Hindi/Urdu/Arabic/Chinese/Russian/Greek is written purely in Latin letters (romanized) without clear, repeated language-specific cue words, return 'ot'. (Only classify as hi/ur/ar/zh/ru/el when native scripts or highly distinctive romanized patterns are clearly present.)

8) Code-heavy inputs:
   - If the text is mostly code with minimal or no natural-language comments/strings, return 'ot'.
   - If comments/strings clearly indicate a language per rules above, use that label.

9) Ambiguity & confidence:
   - When in doubt, choose 'ot' rather than guessing.

Text to classify:
{text}

Output format:
- Respond with EXACTLY one line: "Final Answer: xx"
- Where xx ∈ {{ar, de, el, en, es, fr, hi, ru, tr, ur, vi, zh, ot}} and nothing else.
"""


def parse_final_answer(response: str, debug: bool = False) -> Optional[str]:
    """
    Parse the final answer from the model response.
    For Thinking models, extract from <think>...</think> tags or after them.
    """
    # For Thinking models, the response may have <think></think> tags
    # Remove thinking content and focus on the final answer
    response_stripped = response.strip()
    
    # Remove <think>...</think> content if present
    response_cleaned = re.sub(r'<think>.*?</think>', '', response_stripped, flags=re.DOTALL)
    response_cleaned = response_cleaned.strip()
    
    # Also try the original response
    candidates = [response_cleaned, response_stripped]
    
    valid_labels = {'ar', 'de', 'el', 'en', 'es', 'fr', 'hi', 'ru', 'tr', 'ur', 'vi', 'zh', 'ot'}
    
    # Try multiple patterns to extract language label
    patterns = [
        r"Final Answer:\s*(\w{2})",  # Standard format
        r"Final Answer:\s*([a-z]{2})",  # Lowercase only
        r"Answer:\s*(\w{2})",  # Without "Final"
        r"Language:\s*(\w{2})",  # "Language: xx"
        r"^([a-z]{2})$",  # Just the label alone
        r"\b([a-z]{2})\b\s*$",  # Label at the end with word boundary
        r"is:\s*(\w{2})",  # "is: xx"
        r"→\s*(\w{2})",  # "→ xx"
    ]
    
    for candidate in candidates:
        candidate_lower = candidate.lower()
        
        # Try each pattern
        for pattern in patterns:
            match = re.search(pattern, candidate_lower, re.MULTILINE)
            if match:
                label = match.group(1)
                if label in valid_labels:
                    if debug:
                        print(f"    [DEBUG] Matched pattern '{pattern}' -> '{label}'")
                    return label
        
        # Special case: check if the entire response is just a language code
        if len(candidate) <= 3 and candidate_lower in valid_labels:
            if debug:
                print(f"    [DEBUG] Matched entire response as label -> '{candidate_lower}'")
            return candidate_lower
    
    if debug:
        print(f"    [DEBUG] No pattern matched.")
        print(f"    [DEBUG] Response length: {len(response_stripped)}")
        print(f"    [DEBUG] Cleaned response: '{response_cleaned[:300]}'")
        print(f"    [DEBUG] Original response: '{response_stripped[:300]}'")
    
    return None


async def generate_distillation_data(
    input_file: str,
    output_file: str,
    model_name: str = "Qwen/Qwen3-30B-A3B-Thinking-2507",
    temperature: float = 0.15,
    max_tokens: int = 4096,
    tensor_parallel_size: int = 1,
    max_retries: int = 3,
):
    """
    Generate prompt distillation training data.
    
    Args:
        input_file: Path to file containing sentences to classify (one per line)
        output_file: Path to save the generated training data (JSONL format)
        model_name: Teacher model to use for generating labels
        temperature: Sampling temperature
        max_tokens: Maximum tokens to generate
        tensor_parallel_size: Number of GPUs to use for tensor parallelism
    """
    print(f"Loading input sentences from {input_file}")
    with open(input_file, "r", encoding="utf-8") as f:
        sentences = [line.strip() for line in f if line.strip()]

    print(f"Loaded {len(sentences)} sentences")
    if not sentences:
        print("Input file has no sentences to process, skipping data generation.")
        return

    from vllm import LLM, SamplingParams

    # Initialize vLLM model
    print(f"Initializing teacher model: {model_name}")
    print(f"Using tensor parallelism across {tensor_parallel_size} GPU(s)")
    
    # Get tokenizer to use proper chat template
    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    
    llm = LLM(
        model=model_name,
        tensor_parallel_size=tensor_parallel_size,
        trust_remote_code=True,
        gpu_memory_utilization=0.90,  # Use 90% of GPU memory for better throughput
        max_model_len=32768,  # Match training max length
        enable_prefix_caching=True,  # Cache the system prompt
    )
    
    # Set sampling parameters - use Qwen3 recommended settings
    # For Thinking models, we need to allow enough tokens for reasoning
    sampling_params = SamplingParams(
        temperature=temperature,
        max_tokens=max_tokens,
        top_p=0.8,
        top_k=20,
        # Don't use custom stop sequences - let model finish naturally
        skip_special_tokens=False,  # Keep special tokens for thinking models
    )
    
    # Initial generation
    print("Generating labels with teacher model...")
    results = {}  # sentence -> (response, final_answer)
    failed_indices = []
    failed_examples = []  # Store examples for debugging
    
    # Format prompts using proper chat template
    print("Formatting prompts with Qwen3 chat template...")
    formatted_prompts = []
    for sentence in sentences:
        messages = [
            {
                "role": "user",
                "content": LANGUAGE_CLASSIFICATION_PROMPT.format(text=sentence)
            }
        ]
        # Use tokenizer's chat template
        prompt_text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )
        formatted_prompts.append(prompt_text)
    
    print(f"Sample formatted prompt:")
    print(formatted_prompts[0])
    
    outputs = llm.generate(formatted_prompts, sampling_params)
    
    for idx, (sentence, output) in enumerate(zip(sentences, outputs)):
        response = output.outputs[0].text
        # Enable debug mode for first few failures
        debug_mode = len(failed_examples) < 3
        final_answer = parse_final_answer(response, debug=debug_mode)
        
        if final_answer:
            results[sentence] = (response, final_answer)
        else:
            failed_indices.append(idx)
            # Store first 10 failed examples for debugging
            if len(failed_examples) < 10:
                failed_examples.append({
                    'sentence': sentence,
                    'response': response,
                })
    
    # results is keyed by sentence text, so len(results) counts UNIQUE
    # sentences; the JSONL below writes one row per sentence occurrence. Count
    # rows actually labeled so the reported rate matches the output when the
    # corpus contains duplicate lines (common in language-ID data).
    num_labeled = sum(1 for s in sentences if s in results)
    print(f"\nInitial generation: {num_labeled}/{len(sentences)} successful ({num_labeled/len(sentences)*100:.2f}%)")
    
    # Show debugging info for failed samples
    if failed_examples:
        print(f"\n{'='*60}")
        print("DEBUGGING: Examples of FAILED responses")
        print(f"{'='*60}")
        for i, example in enumerate(failed_examples, 1):
            print(f"\nFailed Example {i}:")
            print(f"  Input: {example['sentence']}")
            print(f"  Response: {example['response']}")
            print(f"  Parsed result: None")
    
    # Show examples of successful responses
    if results:
        print(f"\n{'='*60}")
        print("DEBUGGING: Examples of SUCCESSFUL responses")
        print(f"{'='*60}")
        success_examples = list(results.items())[:3]
        for i, (sentence, (response, label)) in enumerate(success_examples, 1):
            print(f"\nSuccess Example {i}:")
            print(f"  Input: {sentence}")
            print(f"  Response: {response}")
            print(f"  Parsed label: {label}")
    
    # Retry failed generations up to max_retries times
    for retry in range(1, max_retries + 1):
        if not failed_indices:
            break
            
        print(f"\nRetry {retry}/{max_retries}: Regenerating {len(failed_indices)} failed samples...")
        
        # Prepare prompts for failed sentences
        retry_sentences = [sentences[idx] for idx in failed_indices]
        retry_formatted_prompts = []
        for s in retry_sentences:
            messages = [
                {
                    "role": "user",
                    "content": LANGUAGE_CLASSIFICATION_PROMPT.format(text=s)
                }
            ]
            prompt_text = tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True
            )
            retry_formatted_prompts.append(prompt_text)
        
        # Generate with slightly higher temperature to encourage different outputs
        retry_params = SamplingParams(
            temperature=min(temperature * (1 + retry * 0.1), 0.5),  # Gradually increase temp
            max_tokens=max_tokens,
            top_p=0.8,
            top_k=20,
            skip_special_tokens=False,
        )
        
        retry_outputs = llm.generate(retry_formatted_prompts, retry_params)
        
        # Track newly successful and still-failed indices
        new_failed_indices = []
        for idx, sentence, output in zip(failed_indices, retry_sentences, retry_outputs):
            response = output.outputs[0].text
            final_answer = parse_final_answer(response)
            
            if final_answer:
                results[sentence] = (response, final_answer)
            else:
                new_failed_indices.append(idx)
        
        newly_successful = len(failed_indices) - len(new_failed_indices)
        print(f"  ✓ {newly_successful} more samples successful")
        num_labeled = sum(1 for s in sentences if s in results)
        print(f"  Total successful: {num_labeled}/{len(sentences)} ({num_labeled/len(sentences)*100:.2f}%)")
        
        failed_indices = new_failed_indices
    
    # Save results
    print(f"\nSaving results to {output_file}...")
    with open(output_file, "w", encoding="utf-8") as f:
        for sentence in sentences:
            if sentence in results:
                _, final_answer = results[sentence]
                data = {
                    "messages": [
                        {
                            "role": "user",
                            "content": sentence,
                        },
                        {
                            "role": "assistant",
                            "content": final_answer,
                        },
                    ]
                }
                f.write(json.dumps(data, ensure_ascii=False) + "\n")
    
    # Final report
    print(f"\n{'='*60}")
    print("DATA GENERATION COMPLETE")
    print(f"{'='*60}")
    num_labeled = sum(1 for s in sentences if s in results)
    print(f"Total sentences: {len(sentences)}")
    print(f"Valid labels generated: {num_labeled}")
    print(f"Failed after {max_retries} retries: {len(failed_indices)}")
    print(f"Final success rate: {num_labeled/len(sentences)*100:.2f}%")
    print(f"Saved to: {output_file}")
    
    if failed_indices:
        print(f"\n⚠️  Warning: {len(failed_indices)} sentences failed to generate valid labels")
        print("Consider inspecting these samples or adjusting the prompt/temperature")


def main():
    parser = argparse.ArgumentParser(
        description="用教师模型（长提示 + 思考）生成 Prompt 蒸馏训练数据",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--input_file",
        type=str,
        default="./example-data/multilingual.txt",
        help="输入文本文件路径（每行一句待分类文本）",
    )
    parser.add_argument(
        "--output_file",
        type=str,
        default="./data/prompt_distillation_lang.jsonl",
        help="生成的训练数据保存路径（JSONL 格式）",
    )
    parser.add_argument(
        "--model_name",
        type=str,
        default="Qwen/Qwen3-30B-A3B-Thinking-2507",
        help="教师模型名称（用思考型模型以获得更高准确率）",
    )
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.15,
        help="采样温度（与 tinker 保持一致，取 0.15）",
    )
    parser.add_argument(
        "--max_tokens",
        type=int,
        default=4096,
        help="单条生成的最大 token 数",
    )
    parser.add_argument(
        "--tensor_parallel_size",
        type=int,
        default=1,
        help="张量并行使用的 GPU 数（30B 模型在 H100 上建议 2-4）",
    )
    parser.add_argument(
        "--max_retries",
        type=int,
        default=3,
        help="失败样本的最大重试次数",
    )
    
    args = parser.parse_args()
    
    # Create output directory if needed
    output_dir = os.path.dirname(args.output_file)
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)
    
    # Check if input file exists
    if not os.path.exists(args.input_file):
        raise FileNotFoundError(f"Input file not found: {args.input_file}")
    
    # Generate data
    asyncio.run(
        generate_distillation_data(
            input_file=args.input_file,
            output_file=args.output_file,
            model_name=args.model_name,
            temperature=args.temperature,
            max_tokens=args.max_tokens,
            tensor_parallel_size=args.tensor_parallel_size,
            max_retries=args.max_retries,
        )
    )


if __name__ == "__main__":
    main()