| """MiniCPM-V on ZeroGPU: answers a question grounded in the retrieved |
| repair-manual pages. |
| |
| The model and tokenizer are module-level globals: ZeroGPU packs module-level |
| CUDA tensors at startup and shares them with the GPU worker, whereas function |
| arguments are pickled — and trust_remote_code model classes are not picklable. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import spaces |
| import torch |
| from PIL import Image |
| from transformers import AutoModel, AutoTokenizer |
|
|
| from core.constants import ANSWER_GPU_DURATION, ANSWER_MAX_NEW_TOKENS, MINICPM_MODEL_ID |
|
|
| PROMPT = ( |
| "You are a repair-manual assistant. The images are the manual pages most " |
| "relevant to the user's question, each preceded by its label (manual name " |
| "and page number).\n\n" |
| "Answer the question using ONLY what is printed on these pages, following " |
| "these rules:\n" |
| "1. If the answer is a procedure, reproduce EVERY step in order, numbered " |
| "exactly as in the manual. Never skip, merge, or summarize steps. Keep each " |
| "step's notes, model exceptions, and specifications (e.g. torque values) " |
| "with that step, exactly as printed.\n" |
| "2. Quote exact values (torques, clearances, part numbers, capacities) as " |
| "printed, including units.\n" |
| "3. End with the page label(s) you used, e.g. (Manual — p.238). If a " |
| "procedure clearly continues on a page you were not given, say so.\n" |
| "4. If the pages do not contain the answer, say so instead of guessing.\n" |
| "5. Start directly with the answer — no preamble like 'Based on the " |
| "provided pages'.\n\n" |
| "Question: {question}" |
| ) |
|
|
| _MODEL = ( |
| AutoModel.from_pretrained( |
| MINICPM_MODEL_ID, |
| trust_remote_code=True, |
| dtype=torch.bfloat16, |
| attn_implementation="sdpa", |
| ) |
| .to("cuda") |
| .eval() |
| ) |
| _TOKENIZER = AutoTokenizer.from_pretrained(MINICPM_MODEL_ID, trust_remote_code=True) |
|
|
|
|
| @spaces.GPU(duration=ANSWER_GPU_DURATION) |
| def _answer_on_gpu(question: str, pages: list[tuple[str, Image.Image]]) -> str: |
| content = [] |
| for label, img in pages: |
| content.append(f"[{label}]") |
| content.append(img.convert("RGB")) |
| content.append(PROMPT.format(question=question)) |
| with torch.no_grad(): |
| answer = _MODEL.chat( |
| msgs=[{"role": "user", "content": content}], |
| tokenizer=_TOKENIZER, |
| enable_thinking=False, |
| max_new_tokens=ANSWER_MAX_NEW_TOKENS, |
| ) |
| return str(answer).strip() |
|
|
|
|
| class MiniCPM: |
| def answer(self, question: str, pages: list[tuple[str, Image.Image]]) -> str: |
| """pages: [(label, page image)] in retrieval order.""" |
| return _answer_on_gpu(question, pages) |
|
|
