| |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import sys |
| from pathlib import Path |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
|
|
| from scripts.qwen35_qwen3tts_voiceclone_smoke import generate_portuguese_answer, load_llm |
| from speech_bridge_gemma.mimi_overfit import resolve_device, set_seed |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| root = Path(os.environ.get("S2S_ROOT", "/workspace/gemma1b-tts-integration")) |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--root", default=str(root)) |
| parser.add_argument("--device", default="cuda") |
| parser.add_argument("--llm-model", default="Qwen/Qwen3.5-0.8B") |
| parser.add_argument("--llm-dtype", choices=["auto", "float32", "float16", "bfloat16"], default="float16") |
| parser.add_argument("--text-sft-checkpoint", default="job_output/qwen3-text-sft-full-ptbr-20260606T022749Z/qwen_text_sft_layers.pt") |
| parser.add_argument("--system", default="Voce e uma assistente brasileira. Responda sempre em portugues brasileiro, em uma frase curta, direta e natural.") |
| parser.add_argument("--max-answer-tokens", type=int, default=48) |
| parser.add_argument("--temperature", type=float, default=0.0) |
| parser.add_argument("--top-p", type=float, default=0.9) |
| parser.add_argument("--seed", type=int, default=1337) |
| return parser.parse_args() |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
| os.chdir(Path(args.root).resolve()) |
| payload = json.loads(sys.stdin.read() or "{}") |
| question = str(payload.get("question") or "").strip() |
| context = str(payload.get("context") or "").strip() |
| if not question: |
| raise ValueError("question is required") |
| set_seed(args.seed) |
| device = resolve_device(args.device) |
| tokenizer, llm, system, sft = load_llm(args, device) |
| answer = generate_portuguese_answer( |
| llm=llm, |
| tokenizer=tokenizer, |
| question=question, |
| system=system, |
| device=device, |
| max_new_tokens=args.max_answer_tokens, |
| temperature=args.temperature, |
| top_p=args.top_p, |
| context=context, |
| ) |
| print( |
| json.dumps( |
| { |
| "answer": answer, |
| "system": system, |
| "text_sft_loaded_layers": sft.get("text_sft_loaded_layers", 0), |
| }, |
| ensure_ascii=False, |
| ), |
| flush=True, |
| ) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
