Spaces:
Runtime error
Runtime error
| """Generation on the consolidated GPU backend. | |
| Loads Qwen2.5-1.5B-Instruct + the whosouravsharma/paper-qa-lora adapter once | |
| at startup (module-level state, populated by the FastAPI lifespan hook in | |
| app.py). One model instance serves both the fine-tuned behavior (default) | |
| and the raw base model's behavior via peft's disable_adapter() context, so | |
| only one copy of the weights lives in memory. | |
| A single GPU can only run one generate() call at a time, but that must not | |
| block the whole server (health checks, other in-flight requests) while one | |
| generation is running. So actual generation is serialized through a single | |
| asyncio worker task consuming a queue, with the blocking model.generate() | |
| call itself pushed into a one-worker thread executor — the async endpoint | |
| in app.py just awaits its own future. This replaces the earlier archived | |
| design's plain threading.Lock(), which under FastAPI's sync-route threadpool | |
| could let a /health check queue up behind an in-flight generation instead of | |
| returning immediately. | |
| """ | |
| import asyncio | |
| import time | |
| from concurrent.futures import ThreadPoolExecutor | |
| from dataclasses import dataclass | |
| import torch | |
| from peft import PeftModel | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| BASE_MODEL = "Qwen/Qwen2.5-1.5B-Instruct" | |
| ADAPTER_REPO = "whosouravsharma/paper-qa-lora" | |
| def generator_model_name(use_adapter: bool) -> str: | |
| return f"{BASE_MODEL}+LoRA({ADAPTER_REPO})" if use_adapter else BASE_MODEL | |
| SYSTEM_PROMPT = ( | |
| "You are a research paper assistant. Answer the question using only the " | |
| "provided context. Quote the exact sentence(s) you relied on as a " | |
| "citation. If the context does not contain the answer, say so plainly " | |
| "instead of guessing." | |
| ) | |
| state: dict = {} | |
| _executor = ThreadPoolExecutor(max_workers=1) # one GPU: exactly one worker | |
| _queue: "asyncio.Queue | None" = None | |
| _worker_task: "asyncio.Task | None" = None | |
| class GenerateRequest: | |
| context: str | |
| question: str | |
| use_adapter: bool = True | |
| max_new_tokens: int = 256 | |
| def load_model() -> None: | |
| """Called once from app.py's lifespan startup.""" | |
| tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL) | |
| # Qwen2.5's config defaults to bfloat16; force fp16 explicitly rather than | |
| # inheriting that default — this exact mismatch caused real training | |
| # crashes on T4 (see paper-qa-lora/PLAN.md's debugging notes), so it's | |
| # forced here too even though this is inference, not training. | |
| base = AutoModelForCausalLM.from_pretrained(BASE_MODEL, torch_dtype=torch.float16).to("cuda") | |
| model = PeftModel.from_pretrained(base, ADAPTER_REPO) | |
| model.eval() | |
| state["tokenizer"] = tokenizer | |
| state["model"] = model | |
| def _blocking_generate(req: GenerateRequest) -> tuple[str, float]: | |
| tokenizer = state["tokenizer"] | |
| model = state["model"] | |
| messages = [ | |
| {"role": "system", "content": SYSTEM_PROMPT}, | |
| {"role": "user", "content": f"Context:\n{req.context}\n\nQuestion: {req.question}"}, | |
| ] | |
| inputs = tokenizer.apply_chat_template( | |
| messages, add_generation_prompt=True, return_tensors="pt", return_dict=True | |
| ).to(model.device) | |
| start = time.perf_counter() | |
| with torch.no_grad(): | |
| if req.use_adapter: | |
| output = model.generate( | |
| **inputs, max_new_tokens=req.max_new_tokens, do_sample=False, | |
| pad_token_id=tokenizer.eos_token_id, | |
| ) | |
| else: | |
| with model.disable_adapter(): | |
| output = model.generate( | |
| **inputs, max_new_tokens=req.max_new_tokens, do_sample=False, | |
| pad_token_id=tokenizer.eos_token_id, | |
| ) | |
| latency_ms = (time.perf_counter() - start) * 1000 | |
| answer = tokenizer.decode( | |
| output[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True | |
| ).strip() | |
| return answer, latency_ms | |
| async def _worker_loop() -> None: | |
| loop = asyncio.get_running_loop() | |
| while True: | |
| req, future = await _queue.get() | |
| try: | |
| result = await loop.run_in_executor(_executor, _blocking_generate, req) | |
| if not future.done(): | |
| future.set_result(result) | |
| except Exception as e: # noqa: BLE001 - surface any generation failure to the caller | |
| if not future.done(): | |
| future.set_exception(e) | |
| finally: | |
| _queue.task_done() | |
| def start_worker() -> None: | |
| """Called once from app.py's lifespan startup, after load_model().""" | |
| global _queue, _worker_task | |
| _queue = asyncio.Queue() | |
| _worker_task = asyncio.create_task(_worker_loop()) | |
| async def stop_worker() -> None: | |
| if _worker_task is not None: | |
| _worker_task.cancel() | |
| async def generate_answer(req: GenerateRequest) -> tuple[str, float]: | |
| """Enqueues one generation request and awaits its result. Safe to call | |
| concurrently — the queue serializes actual GPU work while this | |
| coroutine itself doesn't block the event loop.""" | |
| future = asyncio.get_running_loop().create_future() | |
| await _queue.put((req, future)) | |
| return await future | |