Update src/llm_handler.py

src/llm_handler.py

@@ -38,9 +38,10 @@ GROQ_MODELS = [
 
 DEFAULT_GROQ_MODEL = "openai/gpt-oss-120b"
 
-DEFAULT_HF_MODEL = "Qwen/Qwen2.5-3B-Instruct"
+DEFAULT_HF_MODEL = "Qwen/Qwen3.5-4B"
 
 HF_MODELS = [
+    "Qwen/Qwen3.5-4B",
     "Qwen/Qwen2.5-3B-Instruct",
     "Qwen/Qwen2.5-1.5B-Instruct",
     "microsoft/Phi-3.5-mini-instruct",
@@ -78,21 +79,48 @@ class HuggingFaceTransformersLLM(LLM):
         dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
         device_map = "auto" if torch.cuda.is_available() else None
 
+        load_kwargs = {
+            "trust_remote_code": True,
+            "device_map": device_map,
+        }
+
         # Prefer `dtype=` (torch_dtype is deprecated in recent transformers)
-        try:
-            self._model = AutoModelForCausalLM.from_pretrained(
-                self.model_name,
-                dtype=dtype,
-                device_map=device_map,
-                trust_remote_code=True,
-            )
-        except TypeError:
-            self._model = AutoModelForCausalLM.from_pretrained(
-                self.model_name,
-                torch_dtype=dtype,
-                device_map=device_map,
-                trust_remote_code=True,
-            )
+        model = None
+        for dtype_key in ("dtype", "torch_dtype"):
+            try:
+                model = AutoModelForCausalLM.from_pretrained(
+                    self.model_name,
+                    **{dtype_key: dtype},
+                    **load_kwargs,
+                )
+                break
+            except TypeError:
+                continue
+            except Exception as causal_err:
+                logger.warning(
+                    "AutoModelForCausalLM failed for %s (%s); trying image-text class",
+                    self.model_name,
+                    causal_err,
+                )
+                break
+
+        if model is None:
+            try:
+                from transformers import AutoModelForImageTextToText
+
+                model = AutoModelForImageTextToText.from_pretrained(
+                    self.model_name,
+                    dtype=dtype,
+                    **load_kwargs,
+                )
+            except Exception:
+                model = AutoModelForCausalLM.from_pretrained(
+                    self.model_name,
+                    torch_dtype=dtype,
+                    **load_kwargs,
+                )
+
+        self._model = model
         if device_map is None:
             self._model = self._model.to("cpu")
         self._model.eval()
@@ -108,23 +136,33 @@ class HuggingFaceTransformersLLM(LLM):
 
         messages = self._build_messages(prompt)
         attention_mask = None
+        enable_thinking = bool(get_config().get("llm.enable_thinking", False))
 
         if hasattr(self._tokenizer, "apply_chat_template"):
-            # Newer transformers may return BatchEncoding when return_dict=True,
-            # or a raw Tensor / BatchEncoding depending on version.
+            template_kwargs: dict[str, Any] = {
+                "add_generation_prompt": True,
+                "return_tensors": "pt",
+                "return_dict": True,
+            }
+            # Qwen3 / Qwen3.5 support enable_thinking in chat template
             try:
                 encoded = self._tokenizer.apply_chat_template(
                     messages,
-                    add_generation_prompt=True,
-                    return_tensors="pt",
-                    return_dict=True,
+                    enable_thinking=enable_thinking,
+                    **template_kwargs,
                 )
             except TypeError:
-                encoded = self._tokenizer.apply_chat_template(
-                    messages,
-                    add_generation_prompt=True,
-                    return_tensors="pt",
-                )
+                try:
+                    encoded = self._tokenizer.apply_chat_template(
+                        messages,
+                        **template_kwargs,
+                    )
+                except TypeError:
+                    encoded = self._tokenizer.apply_chat_template(
+                        messages,
+                        add_generation_prompt=True,
+                        return_tensors="pt",
+                    )
 
             if hasattr(encoded, "input_ids"):
                 input_ids = encoded["input_ids"]