Update app.py

app.py

@@ -1,56 +1,89 @@
 import os
 import gradio as gr
+import onnxruntime as ort
+import numpy as np
+from transformers import AutoTokenizer
 
-from src.model_manager import ModelManager
-from src.inference_engine import InferenceEngine
-
-ASSETS_DIR = "assets"
-MODELS_DIR = os.path.join(ASSETS_DIR, "models")
-
-os.makedirs(ASSETS_DIR, exist_ok=True)
-os.makedirs(MODELS_DIR, exist_ok=True)
-
-manager = ModelManager(MODELS_DIR)
-_ENGINE_CACHE = {}
-
-def list_models():
-    return manager.get_available_models()
-
-def load_engine(model_name: str) -> InferenceEngine:
-    if model_name in _ENGINE_CACHE:
-        return _ENGINE_CACHE[model_name]
-    session, tokenizer, config = manager.load_model(model_name)
-    engine = InferenceEngine(session, tokenizer, config)
-    _ENGINE_CACHE[model_name] = engine
-    return engine
-
-def chat_fn(message, history, model_name, max_tokens, temperature, top_p, top_k):
-    if not model_name:
-        history = history + [{"role": "assistant", "content": "No model selected. Please choose an ONNX model."}]
-        return history
-    try:
-        engine = load_engine(model_name)
-        reply = engine.generate_response(
-            message,
-            max_tokens=int(max_tokens),
-            temperature=float(temperature),
-            top_p=float(top_p),
-            top_k=int(top_k),
-        )
-    except Exception as e:
-        reply = f"Error during inference: {e}"
+ONNX_PATH = os.path.join("assets", "automotive_slm.onnx")
+
+# Load tokenizer (must match training tokenizer)
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# Create ONNX session
+providers = ["CPUExecutionProvider"]
+so = ort.SessionOptions()
+so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+session = ort.InferenceSession(ONNX_PATH, providers=providers, sess_options=so)
+
+# Infer IO names
+INPUT_NAME = session.get_inputs()[0].name
+OUTPUT_NAME = session.get_outputs()[0].name
+
+def generate_onnx(prompt: str, max_tokens=64, temperature=0.8, top_p=0.9, top_k=50) -> str:
+    tokens = tokenizer.encode(prompt)
+    input_ids = np.array([tokens], dtype=np.int64)
+    generated = []
+
+    for _ in range(int(max_tokens)):
+        outputs = session.run([OUTPUT_NAME], {INPUT_NAME: input_ids})
+        logits = outputs[0][0, -1, :]
+
+        # Temperature
+        if temperature and temperature > 0:
+            logits = logits / max(float(temperature), 1e-6)
+
+        # Top-k
+        if top_k and int(top_k) > 0:
+            k = min(int(top_k), logits.shape[-1])
+            idx = np.argpartition(logits, -k)[-k:]
+            filt = np.full_like(logits, -np.inf)
+            filt[idx] = logits[idx]
+            logits = filt
+
+        # Softmax
+        exps = np.exp(logits - np.max(logits))
+        probs = exps / np.sum(exps)
+
+        # Top-p
+        if top_p is not None and 0 < float(top_p) < 1.0:
+            sort_idx = np.argsort(probs)[::-1]
+            sorted_probs = probs[sort_idx]
+            cumsum = np.cumsum(sorted_probs)
+            cutoff = np.searchsorted(cumsum, float(top_p)) + 1
+            mask = np.zeros_like(probs)
+            keep = sort_idx[:cutoff]
+            mask[keep] = probs[keep]
+            s = mask.sum()
+            if s > 0:
+                probs = mask / s
+
+        next_token = int(np.random.choice(len(probs), p=probs))
+        if next_token == tokenizer.eos_token_id:
+            break
+
+        generated.append(next_token)
+        input_ids = np.concatenate([input_ids, [[next_token]]], axis=1)
+
+    text = tokenizer.decode(generated, skip_special_tokens=True).strip()
+    if not text:
+        return "I couldn't generate a response."
+    if text.startswith(prompt):
+        text = text[len(prompt):].strip()
+    return text
+
+def chat_fn(message, history, max_tokens, temperature, top_p, top_k):
+    reply = generate_onnx(message, max_tokens, temperature, top_p, top_k)
     history = history + [
         {"role": "user", "content": message},
         {"role": "assistant", "content": reply},
     ]
     return history
 
-def clear_chat():
-    return []
-
 with gr.Blocks(title="Automotive SLM Chatbot (ONNX)") as demo:
     gr.Markdown("# 🚗 Automotive SLM Chatbot (ONNX-only)")
-    gr.Markdown("Place your .onnx models in assets/models and select one to chat.")
+    gr.Markdown("Using model at assets/automotive_slm.onnx")
 
     with gr.Row():
         with gr.Column(scale=3):
@@ -61,41 +94,23 @@ with gr.Blocks(title="Automotive SLM Chatbot (ONNX)") as demo:
                 clear_btn = gr.Button("Clear")
 
         with gr.Column(scale=2):
-            gr.Markdown("### Model settings")
-            available = list_models()
-            if not available:
-                gr.Markdown("No ONNX models found in assets/models. Please add .onnx files and refresh.")
-                model_dropdown = gr.Dropdown(choices=[], value=None, label="Model", interactive=False)
-                max_tokens = gr.Slider(10, 256, value=64, step=1, label="Max tokens", interactive=False)
-                temperature = gr.Slider(0.1, 1.5, value=0.8, step=0.1, label="Temperature", interactive=False)
-                top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p", interactive=False)
-                top_k = gr.Slider(1, 100, value=50, step=1, label="Top-k", interactive=False)
-            else:
-                # Optional labels with size
-                def size_mb(path):
-                    try: return os.path.getsize(path) / (1024 * 1024)
-                    except Exception: return 0.0
-                labels = [f"{n} ({size_mb(os.path.join(MODELS_DIR, n)):.1f} MB)" for n in available]
-                choices = list(zip(labels, available))
-
-                model_dropdown = gr.Dropdown(choices=choices, value=available[0], label="Model")
-                max_tokens = gr.Slider(10, 256, value=64, step=1, label="Max tokens")
-                temperature = gr.Slider(0.1, 1.5, value=0.8, step=0.1, label="Temperature")
-                top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
-                top_k = gr.Slider(1, 100, value=50, step=1, label="Top-k")
-
-    if available:
-        send_btn.click(
-            fn=chat_fn,
-            inputs=[msg, chatbot, model_dropdown, max_tokens, temperature, top_p, top_k],
-            outputs=[chatbot]
-        )
-        msg.submit(
-            fn=chat_fn,
-            inputs=[msg, chatbot, model_dropdown, max_tokens, temperature, top_p, top_k],
-            outputs=[chatbot]
-        )
-        clear_btn.click(clear_chat, None, chatbot)
+            gr.Markdown("### Generation settings")
+            max_tokens = gr.Slider(10, 256, value=64, step=1, label="Max tokens")
+            temperature = gr.Slider(0.1, 1.5, value=0.8, step=0.1, label="Temperature")
+            top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
+            top_k = gr.Slider(1, 100, value=50, step=1, label="Top-k")
+
+    send_btn.click(
+        fn=chat_fn,
+        inputs=[msg, chatbot, max_tokens, temperature, top_p, top_k],
+        outputs=[chatbot]
+    )
+    msg.submit(
+        fn=chat_fn,
+        inputs=[msg, chatbot, max_tokens, temperature, top_p, top_k],
+        outputs=[chatbot]
+    )
+    clear_btn.click(lambda: [], None, chatbot)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", 7860)))