changed back to transformers

app.py

@@ -1,8 +1,7 @@
 import gradio as gr
 from PIL import Image, ImageDraw
 from ultralytics import YOLO
-from transformers import AutoTokenizer
-from optimum.onnxruntime import ORTModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
 # Load a pre-trained YOLO model
@@ -60,21 +59,16 @@ def get_llm(model_name):
         return llm_cache[model_name]
 
     model_map = {
-        "qwen3:0.6b": "onnx-community/Qwen3-0.6B-ONNX",
-        "gemma3:1b": "onnx-community/gemma-3-1b-it-ONNX-GQA"
-    }
-    hf_model_name = model_map[model_name]
-
-    # Tokenizer is loaded from the original model's repo to ensure correct chat templates
-    original_model_map = {
         "qwen3:0.6b": "Qwen/Qwen3-0.6B-Instruct",
         "gemma3:1b": "google/gemma-3-1b-it"
     }
-    tokenizer = AutoTokenizer.from_pretrained(original_model_map[model_name])
-    model = ORTModelForCausalLM.from_pretrained(
-        hf_model_name, 
-        file_name="model_quantized.onnx",
-        use_cache=False
+    hf_model_name = model_map[model_name]
+
+    tokenizer = AutoTokenizer.from_pretrained(hf_model_name)
+    model = AutoModelForCausalLM.from_pretrained(
+        hf_model_name,
+        torch_dtype=torch.bfloat16,
+        device_map="auto"
     )
 
     llm_cache[model_name] = (model, tokenizer)
@@ -99,20 +93,30 @@ def generate_text(model_name, system_prompt, user_prompt):
         {"role": "user", "content": user_prompt},
     ]
 
-    inputs = tokenizer.apply_chat_template(
+    chat_template_args = {
+        "tokenize": False,
+        "add_generation_prompt": True
+    }
+
+    if 'qwen' in model_name.lower():
+        chat_template_args['enable_thinking'] = False
+
+    text = tokenizer.apply_chat_template(
         messages,
-        add_generation_prompt=True,
-        return_tensors="pt",
+        **chat_template_args
     )
 
-    generated_ids = model.generate(inputs, max_new_tokens=512)
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+    generated_ids = model.generate(
+        model_inputs.input_ids,
+        max_new_tokens=512
+    )
+    generated_ids = [
+        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+    ]
 
     response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    
-    # The response might include the prompt, so we remove it.
-    # This is a common pattern when decoding from a generation.
-    prompt_plus_response = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
-    response = prompt_plus_response[len(tokenizer.decode(inputs[0], skip_special_tokens=True)):]
 
     return response
 

requirements.txt

@@ -3,5 +3,3 @@ ultralytics
 torch
 transformers
 pillow
-bitsandbytes
-optimum[onnxruntime]