| |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import sys |
| from pathlib import Path |
| from typing import Any |
|
|
| import uvicorn |
| from fastapi import FastAPI, Request |
| from fastapi.responses import JSONResponse |
|
|
| 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 |
|
|
|
|
| class AnswerRuntime: |
| def __init__(self, args: argparse.Namespace) -> None: |
| self.args = args |
| self.loaded = False |
|
|
| def load(self) -> None: |
| if self.loaded: |
| return |
| set_seed(self.args.seed) |
| self.device = resolve_device(self.args.device) |
| self.tokenizer, self.llm, self.system, self.sft = load_llm(self.args, self.device) |
| self.loaded = True |
|
|
| def answer(self, payload: dict[str, Any]) -> dict[str, Any]: |
| self.load() |
| question = str(payload.get("question") or "").strip() |
| context = str(payload.get("context") or "").strip() |
| if not question: |
| raise ValueError("question is required") |
| answer = generate_portuguese_answer( |
| llm=self.llm, |
| tokenizer=self.tokenizer, |
| question=question, |
| system=self.system, |
| device=self.device, |
| max_new_tokens=self.args.max_answer_tokens, |
| temperature=self.args.temperature, |
| top_p=self.args.top_p, |
| context=context, |
| ) |
| return { |
| "answer": answer, |
| "system": self.system, |
| "text_sft_loaded_layers": self.sft.get("text_sft_loaded_layers", 0), |
| } |
|
|
|
|
| def build_app(runtime: AnswerRuntime) -> FastAPI: |
| app = FastAPI() |
|
|
| @app.get("/health") |
| async def health() -> dict[str, Any]: |
| return {"ok": True, "loaded": runtime.loaded} |
|
|
| @app.post("/answer") |
| async def answer(request: Request) -> JSONResponse: |
| try: |
| payload = json.loads((await request.body()).decode("utf-8") or "{}") |
| return JSONResponse(runtime.answer(payload)) |
| except Exception as exc: |
| return JSONResponse({"error": str(exc)}, status_code=500) |
|
|
| return app |
|
|
|
|
| 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("--host", default="127.0.0.1") |
| parser.add_argument("--port", type=int, default=8767) |
| 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) |
| args = parser.parse_args() |
| os.chdir(Path(args.root).resolve()) |
| return args |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
| runtime = AnswerRuntime(args) |
| uvicorn.run(build_app(runtime), host=args.host, port=args.port) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
