"""Answer a question about a short passage with a saved ARTI program on CPU.

Run python ask.py --language en (or zh) for the included example.
Use --interactive to paste your own passage and ask several questions.
The package contains public, previously trained adapters; this script does not train.
"""
import argparse
import json
import os
from pathlib import Path
import sys

os.environ.setdefault('TOKENIZERS_PARALLELISM', 'false')
os.environ.setdefault('PYTHONDONTWRITEBYTECODE', '1')
sys.dont_write_bytecode = True
if hasattr(sys.stdout, 'reconfigure'):
    sys.stdout.reconfigure(encoding='utf8')

BASE = 'Qwen/Qwen3-0.6B'
REVISION = 'c1899de289a04d12100db370d81485cdf75e47ca'
ROOT = Path(__file__).resolve().parent
EXAMPLES = {
    'en': {
        'context': 'Orders placed before 2 pm ship the same day. Standard shipping takes 3 business days. Returns are accepted within 30 days of delivery.',
        'question': 'How many days after delivery can an item be returned?',
    },
    'zh': {
        'context': '北城图书馆工作日上午九点开馆，周六上午十点开馆，周日闭馆。图书可以借阅21天，每本图书允许续借一次。',
        'question': '图书可以借阅多久？',
    },
}


def messages(language, context, question):
    instruction = (
        'Read the passage and answer the question. Return only the shortest exact answer from the passage, without explanation.'
        if language == 'en' else '阅读材料并回答问题。只输出材料中最简短、准确的答案，不要解释。'
    )
    return [{'role': 'system', 'content': instruction},
            {'role': 'user', 'content': f'Passage:\n{context}\n\nQuestion: {question}' if language == 'en'
             else f'材料：\n{context}\n\n问题：{question}'}]


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--language', choices=['en', 'zh'], default='en')
    parser.add_argument('--interactive', action='store_true')
    parser.add_argument('--material', type=Path, help='UTF-8 passage file')
    parser.add_argument('--question', help='Question about the passage')
    parser.add_argument('--cache', type=Path, default=ROOT / 'model-cache')
    parser.add_argument('--offline', action='store_true', help='Use already downloaded base weights')
    parser.add_argument('--output', type=Path, default=Path('answers.jsonl'))
    args = parser.parse_args()
    zh = args.language == 'zh'
    say = lambda en, cn: print(cn if zh else en, flush=True)
    import torch
    import arti
    from huggingface_hub import snapshot_download
    from arti.language_recipe import CausalLMPredictor
    if arti.__version__ != '3.3.0':
        raise RuntimeError('This starter targets arti-fit==3.3.0; install requirements.txt.')
    torch.set_num_threads(4)
    say('Loading the public Qwen3 base. The first download is about 1.2 GB; later runs reuse it.',
        '正在加载公开 Qwen3 基座。首次下载约 1.2 GB，后续复用本地文件。')
    base = snapshot_download(BASE, revision=REVISION, cache_dir=args.cache,
                             local_files_only=args.offline,
                             allow_patterns=['*.json', '*.safetensors', '*.jinja', 'LICENSE'])
    say('Base files are available. Loading the model and saved program.',
        '基座文件已准备，正在载入模型和保存的程序。')
    predictor = CausalLMPredictor.from_pretrained(ROOT / 'programs' / args.language,
                    device='cpu', base_model=base, local_files_only=True)
    say('Ready on CPU. No training or GPU is required.', 'CPU 已准备好，无需训练或 GPU。')
    rows = []

    def answer(context, question):
        if not context.strip() or not question.strip():
            raise ValueError('Supply a passage and a question.' if not zh else '请提供材料和问题。')
        turns = messages(args.language, context, question)
        tokens = predictor.tokenizer.apply_chat_template(turns, tokenize=True,
                    add_generation_prompt=True, **predictor.contract['template_kwargs'])
        if len(tokens) + 32 > 384:
            raise ValueError('Shorten the passage: this starter accepts prompt + answer budget <= 384 tokens.'
                             if not zh else '请缩短材料：此入门配置支持材料、问题与答案预算合计不超过 384 tokens。')
        say('Generating an answer on CPU...', '正在 CPU 上生成答案……')
        row = predictor.predict_records([{'id': f'question-{len(rows)+1}', 'messages': turns}],
                                        batch_size=1, max_new_tokens=32)[0]
        observed = {'context': context, 'question': question, 'answer': row['text'],
                    'id': row['id']}
        rows.append(observed)
        print(('问题：' if zh else 'Question: ') + question)
        print(('答案：' if zh else 'Answer: ') + row['text'])
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(''.join(json.dumps(x, ensure_ascii=False)+'\n' for x in rows), encoding='utf8')

    context = args.material.read_text(encoding='utf8') if args.material else EXAMPLES[args.language]['context']
    question = args.question or EXAMPLES[args.language]['question']
    if args.interactive:
        say('Paste a passage. Finish with an empty line. Enter /quit to exit.',
            '粘贴材料，空行结束。输入 /quit 退出。')
        lines = []
        while True:
            line = input()
            if line == '/quit':
                return
            if not line:
                break
            lines.append(line)
        context = '\n'.join(lines)
        while True:
            question = input('问题（/quit 退出）：' if zh else 'Question (/quit to exit): ')
            if question == '/quit':
                break
            answer(context, question)
    else:
        answer(context, question)
    print(('已保存答案到 ' if zh else 'Answers saved to ') + args.output.name)


if __name__ == '__main__':
    main()