fixed qwen, disabled gemma3-1b and minicpm, re-enabled cogito

utils/models.py

@@ -8,24 +8,24 @@ from .prompts import format_rag_prompt
 from .shared import generation_interrupt
 
 models = {
-    "Qwen2.5-1.5b-Instruct": "qwen/qwen2.5-1.5b-instruct",
-    "Qwen2.5-3b-Instruct": "qwen/qwen2.5-3b-instruct",
-    "Llama-3.2-1b-Instruct": "meta-llama/llama-3.2-1b-instruct",
-    "Llama-3.2-3b-Instruct": "meta-llama/llama-3.2-3b-instruct",
-    "Gemma-3-1b-it": "google/gemma-3-1b-it",
+     "Qwen2.5-1.5b-Instruct": "qwen/qwen2.5-1.5b-instruct",
+     "Qwen2.5-3b-Instruct": "qwen/qwen2.5-3b-instruct",
+     "Llama-3.2-1b-Instruct": "meta-llama/llama-3.2-1b-instruct",
+     "Llama-3.2-3b-Instruct": "meta-llama/llama-3.2-3b-instruct",
+    #"Gemma-3-1b-it": "google/gemma-3-1b-it",
     #"Gemma-3-4b-it": "google/gemma-3-4b-it",
-    "Gemma-2-2b-it": "google/gemma-2-2b-it",
-    "Phi-4-mini-instruct": "microsoft/phi-4-mini-instruct",
-    #"Cogito-v1-preview-llama-3b": "deepcogito/cogito-v1-preview-llama-3b",
-    "IBM Granite-3.3-2b-instruct": "ibm-granite/granite-3.3-2b-instruct",
-    #"Bitnet-b1.58-2B4T": "microsoft/bitnet-b1.58-2B-4T",
-    #"MiniCPM3-RAG-LoRA": "openbmb/MiniCPM3-RAG-LoRA",
+     "Gemma-2-2b-it": "google/gemma-2-2b-it",
+     "Phi-4-mini-instruct": "microsoft/phi-4-mini-instruct",
+    "Cogito-v1-preview-llama-3b": "deepcogito/cogito-v1-preview-llama-3b",
+     "IBM Granite-3.3-2b-instruct": "ibm-granite/granite-3.3-2b-instruct",
+    # #"Bitnet-b1.58-2B4T": "microsoft/bitnet-b1.58-2B-4T",
+    # #"MiniCPM3-RAG-LoRA": "openbmb/MiniCPM3-RAG-LoRA",
     "Qwen3-0.6b": "qwen/qwen3-0.6b",
-    "Qwen3-1.7b": "qwen/qwen3-1.7b",
-    "Qwen3-4b": "qwen/qwen3-4b",
-    "SmolLM2-1.7b-Instruct": "HuggingFaceTB/SmolLM2-1.7B-Instruct",
-    "EXAONE-3.5-2.4B-instruct": "LGAI-EXAONE/EXAONE-3.5-2.4B-Instruct",
-    "OLMo-2-1B-Instruct": "allenai/OLMo-2-0425-1B-Instruct",
+     "Qwen3-1.7b": "qwen/qwen3-1.7b",
+     "Qwen3-4b": "qwen/qwen3-4b",
+     "SmolLM2-1.7b-Instruct": "HuggingFaceTB/SmolLM2-1.7B-Instruct",
+     "EXAONE-3.5-2.4B-instruct": "LGAI-EXAONE/EXAONE-3.5-2.4B-Instruct",
+     "OLMo-2-1B-Instruct": "allenai/OLMo-2-0425-1B-Instruct",
 
 }
 
@@ -145,23 +145,28 @@ def run_inference(model_name, context, question):
             device_map='cuda',
             trust_remote_code=True,
             torch_dtype=torch.bfloat16,
+            model_kwargs={
+                "attn_implementation": "eager",
+            }
         )
 
         text_input = format_rag_prompt(question, context, accepts_sys)
-        formatted = tokenizer.apply_chat_template(
-            text_input,
-            tokenize=False,
-            **tokenizer_kwargs,
-        )
+        if "Gemma-3".lower() not in model_name.lower():
+            formatted = pipe.tokenizer.apply_chat_template(
+                text_input,
+                tokenize=False,
+                **tokenizer_kwargs,
+            )
         
-        input_length = len(formatted)
+            input_length = len(formatted)
         # Check interrupt before generation
-        if generation_interrupt.is_set():
-            return ""
 
-        outputs = pipe(formatted, **generation_kwargs)
+            outputs = pipe(formatted, max_new_tokens=512, generation_kwargs={"skip_special_tokens": True})
         #print(outputs[0]['generated_text'])
-        result = outputs[0]['generated_text'][input_length:]
+            result = outputs[0]['generated_text'][input_length:]
+        else: # don't use apply chat template? I don't know why gemma keeps breaking
+            result = pipe(text_input, max_new_tokens=512, generation_kwargs={"skip_special_tokens": True})[0]['generated_text']
+            result = result[0]['generated_text'][-1]['content']
 
     except Exception as e:
         print(f"Error in inference for {model_name}: {e}")