onnx and ui improvements

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from PIL import Image, ImageDraw
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline, AutoTokenizer
+from optimum.onnxruntime import ORTModelForCausalLM
 import torch
 
 # Load the object detection pipeline
@@ -49,27 +50,71 @@ def detect_objects(image):
 
     return annotated_image, detected_objects_str
 
-# Cache for LLM models and tokenizers
+# Cache for LLM models and tokenizers (ONNX Runtime)
 llm_cache = {}
 
-def get_llm(model_name):
-    if model_name in llm_cache:
-        return llm_cache[model_name]
+def get_llm(model_name, preferred_file: str | None = None):
+    cache_key = (model_name, preferred_file or "auto")
+    if cache_key in llm_cache:
+        return llm_cache[cache_key]
 
-    model_map = {
+    # ONNX model repositories on the Hub
+    onnx_repo_map = {
+        "gemma3:1b": "onnx-community/gemma-3-1b-it-ONNX-GQA",
+        "qwen3:0.6b": "onnx-community/Qwen3-0.6B-ONNX",
+    }
+    # Original repos to fetch correct tokenizer + chat templates
+    tokenizer_repo_map = {
         "gemma3:1b": "google/gemma-3-1b-it",
-        "qwen3:0.6b": "Qwen/Qwen3-0.6B-Instruct"
+        "qwen3:0.6b": "Qwen/Qwen3-0.6B-Instruct",
     }
-    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"
-    )
+    onnx_repo = onnx_repo_map[model_name]
+    tokenizer_repo = tokenizer_repo_map[model_name]
+
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_repo)
+
+    # Try a few common ONNX filenames found in community repos to avoid the
+    # "Too many ONNX model files were found" ambiguity.
+    candidate_files = [
+        "model_q4.onnx",
+        "model_quantized.onnx",
+        "model_int8.onnx",
+        "model.onnx",
+    ]
 
-    llm_cache[model_name] = (model, tokenizer)
+    model = None
+    last_err = None
+    ordered = candidate_files
+    if preferred_file and preferred_file in candidate_files:
+        # Put preferred file first
+        ordered = [preferred_file] + [f for f in candidate_files if f != preferred_file]
+    elif preferred_file and preferred_file not in candidate_files:
+        # If user typed a specific known filename not in our shortlist, try it first anyway
+        ordered = [preferred_file] + candidate_files
+
+    for fname in ordered:
+        try:
+            model = ORTModelForCausalLM.from_pretrained(
+                onnx_repo,
+                subfolder="onnx",
+                file_name=fname,
+            )
+            break
+        except Exception as e:
+            last_err = e
+            continue
+    if model is None:
+        raise RuntimeError(f"Failed to load ONNX model from {onnx_repo}. Last error: {last_err}")
+
+    # Disable cache to avoid past_key_values shape issues on some ONNX builds
+    if hasattr(model.config, "use_cache"):
+        try:
+            model.config.use_cache = False
+        except Exception:
+            pass
+
+    llm_cache[cache_key] = (model, tokenizer)
     return model, tokenizer
 
 def update_user_prompt(detected_objects, current_prompt):
@@ -83,40 +128,115 @@ def update_user_prompt(detected_objects, current_prompt):
         
     return new_prompt
 
-def generate_text(model_name, system_prompt, user_prompt):
-    model, tokenizer = get_llm(model_name)
+def generate_text(
+    model_name,
+    onnx_file_choice,
+    system_prompt,
+    user_prompt,
+    do_sample,
+    temperature,
+    top_p,
+    top_k,
+    repetition_penalty,
+    max_new_tokens,
+):
+    model, tokenizer = get_llm(model_name, preferred_file=None if onnx_file_choice == "auto" else onnx_file_choice)
 
     messages = [
         {"role": "system", "content": system_prompt},
         {"role": "user", "content": user_prompt},
     ]
 
-    chat_template_args = {
+    chat_template_kwargs = {
         "tokenize": False,
-        "add_generation_prompt": True
+        "add_generation_prompt": True,
     }
-
-    if 'qwen' in model_name.lower():
-        chat_template_args['enable_thinking'] = False
+    # Disable "thinking" for Qwen models
+    if "qwen" in model_name.lower():
+        chat_template_kwargs["enable_thinking"] = False
 
     text = tokenizer.apply_chat_template(
         messages,
-        **chat_template_args
+        **chat_template_kwargs,
     )
 
-    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)
+    inputs = tokenizer([text], return_tensors="pt")
+
+    with torch.inference_mode():
+        gen_ids = model.generate(
+            **inputs,
+            max_new_tokens=int(max_new_tokens),
+            do_sample=bool(do_sample),
+            temperature=float(temperature),
+            top_p=float(top_p),
+            top_k=int(top_k),
+            repetition_penalty=float(repetition_penalty),
+        )
+
+    # Decode only the newly generated tokens beyond the input length
+    trimmed = [
+        output_ids[len(input_ids):]
+        for input_ids, output_ids in zip(inputs.input_ids, gen_ids)
     ]
+    response = tokenizer.batch_decode(trimmed, skip_special_tokens=True)[0]
+    return response
 
-    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+def chat_respond(
+    model_name,
+    onnx_file_choice,
+    system_prompt,
+    message,
+    history,
+    do_sample,
+    temperature,
+    top_p,
+    top_k,
+    repetition_penalty,
+    max_new_tokens,
+):
+    """Builds a chat messages list from history + current user message, generates a reply, and returns updated history and an empty input box."""
+    # Guard: empty message
+    if not (message and message.strip()):
+        return history, gr.update(value="")
+
+    # Build messages: system, then alternating user/assistant from history, then current user
+    messages = [{"role": "system", "content": system_prompt}]
+    for u, a in (history or []):
+        if u:
+            messages.append({"role": "user", "content": u})
+        if a:
+            messages.append({"role": "assistant", "content": a})
+    messages.append({"role": "user", "content": message})
+
+    # Generate using the same path as generate_text, but inline to avoid extra serialization
+    model, tokenizer = get_llm(model_name, preferred_file=None if onnx_file_choice == "auto" else onnx_file_choice)
+
+    chat_template_kwargs = {
+        "tokenize": False,
+        "add_generation_prompt": True,
+    }
+    if "qwen" in model_name.lower():
+        chat_template_kwargs["enable_thinking"] = False
 
-    return response
+    text = tokenizer.apply_chat_template(messages, **chat_template_kwargs)
+    inputs = tokenizer([text], return_tensors="pt")
+
+    with torch.inference_mode():
+        gen_ids = model.generate(
+            **inputs,
+            max_new_tokens=int(max_new_tokens),
+            do_sample=bool(do_sample),
+            temperature=float(temperature),
+            top_p=float(top_p),
+            top_k=int(top_k),
+            repetition_penalty=float(repetition_penalty),
+        )
+
+    trimmed = [output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, gen_ids)]
+    reply = tokenizer.batch_decode(trimmed, skip_special_tokens=True)[0]
+
+    new_history = (history or []) + [(message, reply)]
+    return new_history, gr.update(value="")
 
 with gr.Blocks() as demo:
     gr.Markdown("# Black Box: Object Detection and LLM Chat")
@@ -130,10 +250,25 @@ with gr.Blocks() as demo:
 
     with gr.Tab("LLM Chat"):
         model_selector = gr.Dropdown(choices=["gemma3:1b", "qwen3:0.6b"], label="Select LLM Model")
+        onnx_file_selector = gr.Dropdown(
+            choices=["auto", "model_q4.onnx", "model_int8.onnx", "model_quantized.onnx", "model.onnx"],
+            value="auto",
+            label="ONNX file variant"
+        )
         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")
-        llm_button = gr.Button("Generate")
+        chat_bot = gr.Chatbot(height=360, label="Conversation")
+        chat_history = gr.State([])
+        user_prompt_input = gr.Textbox(label="Message", placeholder="Type your message and press Send...", lines=3)
+        with gr.Accordion("Generation settings", open=False):
+            do_sample_cb = gr.Checkbox(value=True, label="do_sample")
+            temperature_sl = gr.Slider(minimum=0.0, maximum=2.0, value=0.7, step=0.05, label="temperature")
+            top_p_sl = gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.01, label="top_p")
+            top_k_sl = gr.Slider(minimum=0, maximum=200, value=50, step=1, label="top_k")
+            repetition_penalty_sl = gr.Slider(minimum=0.8, maximum=2.0, value=1.05, step=0.01, label="repetition_penalty")
+            max_new_tokens_sl = gr.Slider(minimum=1, maximum=1024, value=512, step=1, label="max_new_tokens")
+        with gr.Row():
+            send_btn = gr.Button("Send", variant="primary")
+            clear_btn = gr.Button("Clear chat")
 
     # Connect object detection components
     object_detection_button.click(
@@ -142,18 +277,52 @@ with gr.Blocks() as demo:
         outputs=[detected_image_output, detected_objects_output]
     )
 
-    # Connect LLM components
-    llm_button.click(
-        fn=generate_text,
-        inputs=[model_selector, system_prompt_input, user_prompt_input],
-        outputs=llm_output
+    # Connect LLM chat components
+    send_btn.click(
+        fn=chat_respond,
+        inputs=[
+            model_selector,
+            onnx_file_selector,
+            system_prompt_input,
+            user_prompt_input,
+            chat_history,
+            do_sample_cb,
+            temperature_sl,
+            top_p_sl,
+            top_k_sl,
+            repetition_penalty_sl,
+            max_new_tokens_sl,
+        ],
+        outputs=[chat_bot, user_prompt_input],
+    )
+    # Also submit on Enter
+    user_prompt_input.submit(
+        fn=chat_respond,
+        inputs=[
+            model_selector,
+            onnx_file_selector,
+            system_prompt_input,
+            user_prompt_input,
+            chat_history,
+            do_sample_cb,
+            temperature_sl,
+            top_p_sl,
+            top_k_sl,
+            repetition_penalty_sl,
+            max_new_tokens_sl,
+        ],
+        outputs=[chat_bot, user_prompt_input],
     )
+    # Clear chat
+    def _clear_chat():
+        return [], gr.update(value="")
+    clear_btn.click(fn=_clear_chat, inputs=None, outputs=[chat_bot, user_prompt_input])
 
-    # Connect detected objects to user prompt
+    # Connect detected objects to user message input
     detected_objects_output.change(
         fn=update_user_prompt,
         inputs=[detected_objects_output, user_prompt_input],
-        outputs=user_prompt_input
+        outputs=user_prompt_input,
     )
 
 demo.launch()

requirements.txt

@@ -2,4 +2,5 @@ gradio
 torch
 transformers
 pillow
-accelerate
+accelerate
+optimum[onnxruntime]