"""Dual-model inference. Construct the backend at application module scope, before launch.""" from __future__ import annotations import json import os import re import time from .config import SETTINGS from . import prompts from .validation import json_object, OutputError _ACTIVE_BACKEND = None _GPU_TASK = None def _gpu_duration(task: str, payload_json: str) -> int: return {"question": SETTINGS.question_seconds, "grade": SETTINGS.grading_seconds, "explain": SETTINGS.hint_seconds}[task] def _gpu_entry(task: str, payload_json: str) -> str: # Module-level callable, as in the original working ZeroGPU pattern. if _ACTIVE_BACKEND is None: raise RuntimeError("Inference backend was not initialized at startup.") return _ACTIVE_BACKEND._dispatch(task, payload_json) class TransformersBackend: def __init__(self): # Must precede torch/transformers (app.py imports spaces before Gradio too). import spaces import torch from transformers import AutoModelForCausalLM, AutoTokenizer self.torch = torch self.models = {} self.device = os.getenv("LLM_DEVICE", "cuda") if self.device == "cuda" and torch.version.cuda is None: raise RuntimeError("A CUDA-enabled PyTorch runtime is required. Use DEMO_MODE=1 for CPU-only UI tests.") dtype = torch.bfloat16 if self.device == "cuda" else torch.float32 for name in dict.fromkeys([SETTINGS.question_model, SETTINGS.grading_model]): tokenizer = AutoTokenizer.from_pretrained(name, trust_remote_code=False) model = AutoModelForCausalLM.from_pretrained(name, torch_dtype=dtype, low_cpu_mem_usage=True, trust_remote_code=False) model.to(self.device) model.eval() self.models[name] = (model, tokenizer) print(f"[Models] question={SETTINGS.question_model}; evaluator={SETTINGS.grading_model}; device={self.device}", flush=True) global _ACTIVE_BACKEND, _GPU_TASK _ACTIVE_BACKEND = self _GPU_TASK = spaces.GPU(duration=_gpu_duration, size="large")(_gpu_entry) def _dispatch(self, task: str, payload_json: str) -> str: print(f"[Inference worker] task={task} device={self.device}", flush=True) payload = json.loads(payload_json) if task == "question": draft = self._generate(prompts.QUESTION, payload, SETTINGS.question_model, 1600) from .schemas import Chunk from .validation import validate_draft, apply_audit source_objects = [Chunk(**c) for c in payload["sources"]] parsed = validate_draft(draft, source_objects) audit = self._generate(prompts.AUDIT, {"sources": payload["sources"], "candidate": parsed.model_dump()}, SETTINGS.question_model, 420) final = apply_audit(parsed, audit) print(f"[Rubric audit] draft_items={len(parsed.concepts)} final_items={len(final.concepts)}", flush=True) return final.model_dump_json() if task == "grade": return json.dumps(self._generate(prompts.ASSESS, payload, SETTINGS.grading_model, 1800)) if task == "explain": return json.dumps(self._generate(prompts.EXPLAIN, payload, SETTINGS.question_model, 650)) raise ValueError("Unknown task.") def _generate(self, system: str, data: dict, model_id: str, new_tokens: int) -> dict: model, tokenizer = self.models[model_id] text = tokenizer.apply_chat_template([{"role": "system", "content": system}, {"role": "user", "content": json.dumps(data, ensure_ascii=False)}], tokenize=False, add_generation_prompt=True, enable_thinking=False) inputs = tokenizer(text, return_tensors="pt") count = int(inputs.input_ids.shape[1]) if count > SETTINGS.max_prompt_tokens: raise OutputError("This question has too much context. Choose a narrower topic or shorten the answer.") inputs = inputs.to(self.device) print(f"[LLM] model={model_id} input_tokens={count} output_limit={new_tokens}", flush=True) started = time.monotonic() with self.torch.inference_mode(): output = model.generate(**inputs, max_new_tokens=new_tokens, do_sample=False, pad_token_id=tokenizer.eos_token_id, use_cache=True) generated = output[0, count:] raw = tokenizer.decode(generated, skip_special_tokens=True) print(f"[LLM] generated_tokens={len(generated)} seconds={time.monotonic()-started:.1f}", flush=True) # Decode and validate in the GPU worker; only plain JSON crosses back. return json_object(raw) def call(self, task: str, payload: dict) -> dict: return json.loads(_GPU_TASK(task, json.dumps(payload, ensure_ascii=False))) class DemoBackend: """Deterministic fixtures for local UI/testing, prominently labelled. Not a tutor.""" def call(self, task: str, payload: dict) -> dict: if task == "question": sources = payload["sources"][:2] aspects = ["main concepts"] concepts = [] for i, c in enumerate(sources, 1): quote = c["text"][:200].strip() concepts.append({"id": f"C{i}", "name": f"Source concept {i}", "description": quote, "importance": "essential" if i == 1 else "important", "aspect": aspects[0], "relevance": "The question asks for the main concepts in these passages.", "evidence": [{"source_id": c["id"], "quote": quote}]}) return {"question": f"Describe the main concepts in {payload['topic_label']}.", "asked_aspects": aspects, "concepts": concepts} if task == "grade": checks = [] answer = payload["answer"] for c in payload["rubric"]["concepts"]: covered = c["description"].casefold() in answer.casefold() checks.append({"concept_id": c["id"], "status": "covered" if covered else "missing", "answer_evidence": c["description"] if covered else "", "missing_detail": "Demo fixture: this source phrase was not reproduced.", "what_to_add": c["description"]}) return {"checks": checks} return {"notes": []} def create_backend(): return DemoBackend() if SETTINGS.demo else TransformersBackend() class LocalCPUBackend(TransformersBackend): """Same constrained task pipeline as ZeroGPU, with a CPU model cache. Constructed in a fresh CPU subprocess, not in the hosted ZeroGPU parent. Only one model is kept resident at a time to bound memory use. """ def __init__(self): import torch self.torch = torch self.device = 'cpu' self.models = {} torch.set_num_threads(int(os.getenv('LOCAL_CPU_THREADS', '4'))) def _generate(self, system, data, model_id, new_tokens): if model_id not in self.models: import gc self.models.clear() gc.collect() from transformers import AutoModelForCausalLM, AutoTokenizer print(f"[CPU LLM] loading {model_id}; no remote inference", flush=True) offline = os.getenv('HF_HUB_OFFLINE', '0') == '1' tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=False, token=False, local_files_only=offline) model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=self.torch.float32, low_cpu_mem_usage=True, trust_remote_code=False, token=False, local_files_only=offline) model.to('cpu') model.eval() self.models[model_id] = (model, tokenizer) return super()._generate(system, data, model_id, new_tokens)