Spaces:
Runtime error
Runtime error
| from __future__ import annotations | |
| import json | |
| import re | |
| import threading | |
| from dataclasses import dataclass | |
| from typing import Any | |
| from PIL import Image | |
| from .config import SETTINGS | |
| from .normalization import enrich_characteristic, records_from_ocr_text | |
| from .schema import Category, Characteristic, ReviewStatus | |
| class BackendResponse: | |
| characteristics: list[Characteristic] | |
| drawing_metadata: dict[str, Any] | |
| general_notes: list[str] | |
| raw_text: str | |
| class QwenBackend: | |
| def __init__(self, model_id: str | None = None) -> None: | |
| self.model_id = model_id or SETTINGS.qwen_model_id | |
| self._model: Any = None | |
| self._processor: Any = None | |
| self._lock = threading.Lock() | |
| def _load(self) -> None: | |
| if self._model is not None: | |
| return | |
| with self._lock: | |
| if self._model is not None: | |
| return | |
| try: | |
| import torch | |
| from transformers import AutoModelForImageTextToText, AutoProcessor | |
| except ImportError as exc: | |
| raise RuntimeError( | |
| "Qwen dependencies are unavailable. Install torch and transformers " | |
| "or use Demo mode." | |
| ) from exc | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| if device == "cuda": | |
| capability = torch.cuda.get_device_capability() | |
| dtype = torch.bfloat16 if capability[0] >= 8 else torch.float16 | |
| else: | |
| dtype = torch.float32 | |
| kwargs: dict[str, Any] = { | |
| "torch_dtype": dtype, | |
| "low_cpu_mem_usage": True, | |
| "token": SETTINGS.hf_token, | |
| } | |
| if device == "cuda": | |
| kwargs["device_map"] = "auto" | |
| kwargs["attn_implementation"] = "sdpa" | |
| self._processor = AutoProcessor.from_pretrained( | |
| self.model_id, token=SETTINGS.hf_token | |
| ) | |
| self._model = AutoModelForImageTextToText.from_pretrained( | |
| self.model_id, **kwargs | |
| ) | |
| if device == "cpu": | |
| self._model.to(device) | |
| self._model.eval() | |
| def extract(self, image: Image.Image, prompt: str, page: int) -> BackendResponse: | |
| self._load() | |
| import torch | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "image", "image": image.convert("RGB")}, | |
| {"type": "text", "text": prompt}, | |
| ], | |
| } | |
| ] | |
| inputs = self._processor.apply_chat_template( | |
| messages, | |
| tokenize=True, | |
| add_generation_prompt=True, | |
| return_dict=True, | |
| return_tensors="pt", | |
| ) | |
| inputs = inputs.to(self._model.device) | |
| with torch.inference_mode(): | |
| generated = self._model.generate( | |
| **inputs, | |
| max_new_tokens=SETTINGS.max_new_tokens, | |
| do_sample=False, | |
| use_cache=True, | |
| ) | |
| prompt_length = inputs["input_ids"].shape[-1] | |
| text = self._processor.decode( | |
| generated[0][prompt_length:], skip_special_tokens=True | |
| ).strip() | |
| payload = extract_json_payload(text) | |
| return response_from_payload(payload, text, page) | |
| class PaddleBackend: | |
| """Experimental OCR baseline; it intentionally does not claim semantic accuracy.""" | |
| def __init__(self, model_id: str | None = None) -> None: | |
| self.model_id = model_id or SETTINGS.paddle_model_id | |
| self._model: Any = None | |
| self._processor: Any = None | |
| self._lock = threading.Lock() | |
| def _load(self) -> None: | |
| if self._model is not None: | |
| return | |
| with self._lock: | |
| if self._model is not None: | |
| return | |
| try: | |
| import torch | |
| from transformers import AutoModelForImageTextToText, AutoProcessor | |
| except ImportError as exc: | |
| raise RuntimeError( | |
| "PaddleOCR-VL dependencies are unavailable. Install torch and " | |
| "transformers>=5 or use Demo mode." | |
| ) from exc | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| if device == "cuda": | |
| capability = torch.cuda.get_device_capability() | |
| dtype = torch.bfloat16 if capability[0] >= 8 else torch.float16 | |
| else: | |
| dtype = torch.float32 | |
| self._processor = AutoProcessor.from_pretrained( | |
| self.model_id, token=SETTINGS.hf_token, trust_remote_code=True | |
| ) | |
| kwargs: dict[str, Any] = { | |
| "torch_dtype": dtype, | |
| "low_cpu_mem_usage": True, | |
| "token": SETTINGS.hf_token, | |
| "trust_remote_code": True, | |
| } | |
| if device == "cuda": | |
| kwargs["device_map"] = "auto" | |
| self._model = AutoModelForImageTextToText.from_pretrained( | |
| self.model_id, **kwargs | |
| ) | |
| if device == "cpu": | |
| self._model.to(device) | |
| self._model.eval() | |
| def recognize(self, image: Image.Image, page: int) -> BackendResponse: | |
| self._load() | |
| import torch | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "image", "image": image.convert("RGB")}, | |
| {"type": "text", "text": "OCR:"}, | |
| ], | |
| } | |
| ] | |
| inputs = self._processor.apply_chat_template( | |
| messages, | |
| add_generation_prompt=True, | |
| tokenize=True, | |
| return_dict=True, | |
| return_tensors="pt", | |
| ).to(self._model.device) | |
| with torch.inference_mode(): | |
| output = self._model.generate( | |
| **inputs, | |
| max_new_tokens=SETTINGS.max_new_tokens, | |
| do_sample=False, | |
| ) | |
| prompt_length = inputs["input_ids"].shape[-1] | |
| text = self._processor.decode( | |
| output[0][prompt_length:], skip_special_tokens=True | |
| ).strip() | |
| return BackendResponse(records_from_ocr_text(text, page), {}, [], text) | |
| def extract_json_payload(text: str) -> dict[str, Any]: | |
| cleaned = text.strip() | |
| cleaned = re.sub(r"^```(?:json)?\s*", "", cleaned, flags=re.I) | |
| cleaned = re.sub(r"\s*```$", "", cleaned) | |
| decoder = json.JSONDecoder() | |
| starts = [index for index in (cleaned.find("{"), cleaned.find("[")) if index >= 0] | |
| if not starts: | |
| raise ValueError("The model response did not contain JSON.") | |
| start = min(starts) | |
| try: | |
| payload, _ = decoder.raw_decode(cleaned[start:]) | |
| except json.JSONDecodeError as exc: | |
| raise ValueError(f"The model returned malformed JSON: {exc.msg}.") from exc | |
| if isinstance(payload, list): | |
| return {"characteristics": payload} | |
| if not isinstance(payload, dict): | |
| raise ValueError("The model JSON must be an object or a list of characteristics.") | |
| return payload | |
| def response_from_payload( | |
| payload: dict[str, Any], raw_text: str, page: int | |
| ) -> BackendResponse: | |
| raw_records = payload.get("characteristics", payload.get("parameters", [])) | |
| if not isinstance(raw_records, list): | |
| raw_records = [] | |
| records: list[Characteristic] = [] | |
| for raw in raw_records: | |
| if not isinstance(raw, dict): | |
| continue | |
| normalized = _normalize_record_dict(raw, page) | |
| try: | |
| records.append(enrich_characteristic(Characteristic.model_validate(normalized))) | |
| except Exception: | |
| # Preserve partially useful content instead of failing a complete drawing. | |
| fallback = Characteristic( | |
| balloon_id=str(raw.get("balloon_id", raw.get("balloon", ""))), | |
| page=page, | |
| requirement=str(raw.get("requirement", raw.get("text", ""))), | |
| source_text=str(raw.get("source_text", raw.get("raw_text", ""))), | |
| category=Category.OTHER, | |
| confidence=0.2, | |
| status=ReviewStatus.NEEDS_REVIEW, | |
| validation_issues=["Model record required schema recovery."], | |
| ) | |
| records.append(enrich_characteristic(fallback)) | |
| metadata = payload.get("drawing_metadata", {}) | |
| if not isinstance(metadata, dict): | |
| metadata = {} | |
| notes = payload.get("general_notes", []) | |
| if isinstance(notes, str): | |
| notes = [notes] | |
| notes = [str(note).strip() for note in notes if str(note).strip()] | |
| return BackendResponse(records, metadata, notes, raw_text) | |
| def _normalize_record_dict(raw: dict[str, Any], page: int) -> dict[str, Any]: | |
| item = dict(raw) | |
| aliases = { | |
| "balloon": "balloon_id", | |
| "balloon_number": "balloon_id", | |
| "type": "category", | |
| "parameter_type": "category", | |
| "raw_text": "source_text", | |
| "bbox": "source_bbox", | |
| "center": "balloon_center", | |
| "upper_tol": "upper_tolerance", | |
| "lower_tol": "lower_tolerance", | |
| "primary_datum": "datum_primary", | |
| "secondary_datum": "datum_secondary", | |
| "tertiary_datum": "datum_tertiary", | |
| } | |
| for old, new in aliases.items(): | |
| if new not in item and old in item: | |
| item[new] = item[old] | |
| item["page"] = _integer(item.get("page"), page) | |
| item["category"] = _category(item.get("category", "Other")) | |
| item["status"] = ReviewStatus.NEEDS_REVIEW | |
| return item | |
| def _category(value: Any) -> str: | |
| text = re.sub(r"[^a-z]+", " ", str(value).casefold()).strip() | |
| mappings = { | |
| "linear dimension": Category.LINEAR_DIMENSION, | |
| "dimension": Category.LINEAR_DIMENSION, | |
| "measure": Category.LINEAR_DIMENSION, | |
| "diameter": Category.DIAMETER, | |
| "radius": Category.RADIUS, | |
| "angle": Category.ANGLE, | |
| "gd t": Category.GD_AND_T, | |
| "gdt": Category.GD_AND_T, | |
| "geometric tolerance": Category.GD_AND_T, | |
| "thread": Category.THREAD, | |
| "surface finish": Category.SURFACE_FINISH, | |
| "surface roughness": Category.SURFACE_FINISH, | |
| "general tolerance": Category.GENERAL_TOLERANCE, | |
| "material": Category.MATERIAL, | |
| "note": Category.NOTE, | |
| } | |
| return mappings.get(text, Category.OTHER).value | |
| def _integer(value: Any, default: int) -> int: | |
| try: | |
| return int(value) | |
| except (TypeError, ValueError): | |
| return default | |
| class BackendRegistry: | |
| """Process-local model cache; Spaces normally run one worker per GPU.""" | |
| def __init__(self) -> None: | |
| self._qwen: QwenBackend | None = None | |
| self._paddle: PaddleBackend | None = None | |
| def qwen(self) -> QwenBackend: | |
| if self._qwen is None: | |
| self._qwen = QwenBackend() | |
| return self._qwen | |
| def paddle(self) -> PaddleBackend: | |
| if self._paddle is None: | |
| self._paddle = PaddleBackend() | |
| return self._paddle | |
| BACKENDS = BackendRegistry() | |