test

app.py

@@ -1,8 +1,9 @@
 import gradio as gr
 from PIL import Image, ImageDraw
 from ultralytics import YOLO
+from transformers import AutoTokenizer
+from optimum.onnxruntime import ORTModelForCausalLM
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # Load a pre-trained YOLO model
 model = YOLO('yolov8n.pt')
@@ -59,13 +60,18 @@ def get_llm(model_name):
         return llm_cache[model_name]
 
     model_map = {
-        "qwen3:0.6b": "Qwen/Qwen3-0.6B-Instruct",
-        "gemma3:1b": "google/gemma-3-1b-it"
+        "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 = AutoTokenizer.from_pretrained(hf_model_name)
-    model = AutoModelForCausalLM.from_pretrained(hf_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)
 
     llm_cache[model_name] = (model, tokenizer)
     return model, tokenizer
@@ -88,24 +94,21 @@ def generate_text(model_name, system_prompt, user_prompt):
         {"role": "system", "content": system_prompt},
         {"role": "user", "content": user_prompt},
     ]
-    
-    text = tokenizer.apply_chat_template(
+
+    inputs = tokenizer.apply_chat_template(
         messages,
-        tokenize=False,
-        add_generation_prompt=True
+        add_generation_prompt=True,
+        return_tensors="pt",
     )
-    
-    model_inputs = tokenizer([text], return_tensors="pt")
 
-    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)
-    ]
+    generated_ids = model.generate(inputs, max_new_tokens=512)
 
     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
 
@@ -120,7 +123,7 @@ with gr.Blocks() as demo:
         detected_objects_output = gr.Textbox(label="Detected Objects")
 
     with gr.Tab("LLM Chat"):
-        model_selector = gr.Dropdown(choices=["qwen2:0.5b", "gemma2:2b"], label="Select LLM Model")
+        model_selector = gr.Dropdown(choices=["qwen3:0.6b", "gemma3:1b"], label="Select LLM Model")
         system_prompt_input = gr.Textbox(label="System Prompt", value="You are a helpful assistant.")
         user_prompt_input = gr.Textbox(label="User Prompt")
         llm_output = gr.Textbox(label="LLM Response")

requirements.txt

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