captcha_solver/engines/qwen_engine.py

"""Qwen2.5 small text LLM engine for math / reasoning / OCR cleanup."""

from __future__ import annotations

import os
from typing import Optional

from captcha_solver.engines.base import BaseEngine
from captcha_solver.config import get_settings


class QwenEngine(BaseEngine):
    name = "qwen"

    def __init__(self) -> None:
        super().__init__()
        self._model = None
        self._tokenizer = None

    def _do_load(self) -> None:
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer

        s = get_settings()
        os.environ.setdefault("HF_HOME", str(s.cache_dir / "hf"))
        dtype = {"float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16}[
            s.qwen_torch_dtype
        ]

        self._tokenizer = AutoTokenizer.from_pretrained(s.qwen_model)
        self._model = AutoModelForCausalLM.from_pretrained(
            s.qwen_model,
            torch_dtype=dtype,
            low_cpu_mem_usage=True,
        ).to(s.qwen_device).eval()

    def _do_unload(self) -> None:
        self._model = None
        self._tokenizer = None

    def generate(
        self,
        prompt: str,
        system: Optional[str] = None,
        max_new_tokens: int = 64,
        temperature: float = 0.0,
    ) -> str:
        """Generate a short response. Best for math, OCR cleanup, simple reasoning."""
        if not self._loaded:
            self.load()
        import torch
        assert self._model is not None and self._tokenizer is not None

        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})
        text = self._tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        inputs = self._tokenizer(text, return_tensors="pt").to(self._model.device)
        with torch.no_grad():
            gen = self._model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=temperature > 0,
                temperature=temperature or 1.0,
                pad_token_id=self._tokenizer.eos_token_id,
            )
        out = gen[0][inputs["input_ids"].shape[1]:]
        return self._tokenizer.decode(out, skip_special_tokens=True).strip()