engine/llm.py

"""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)