Spaces:
Runtime error
Runtime error
File size: 5,117 Bytes
96c78fb 3873f6a 96c78fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | """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
@dataclass
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
|