Update app.py

app.py

@@ -1,87 +1,143 @@
 import os
 import gradio as gr
+
 from src.model_manager import ModelManager
 from src.inference_engine import InferenceEngine
 
 ASSETS_DIR = "assets"
+MODELS_DIR = os.path.join(ASSETS_DIR, "models")
 
-# Initialize once
-manager = ModelManager(os.path.join(ASSETS_DIR, "models"))
+# Ensure directories exist (prevents path issues)
+os.makedirs(ASSETS_DIR, exist_ok=True)
+os.makedirs(MODELS_DIR, exist_ok=True)
 
-def list_models():
-    models = manager.get_available_models()
-    return models
+# Initialize global model manager
+manager = ModelManager(MODELS_DIR)
 
-# Cache loaded engines by model name
-_engines = {}
+# Cache of InferenceEngine per model filename
+_ENGINE_CACHE = {}
+
+def list_models():
+    """Return available model filenames from assets/models"""
+    return manager.get_available_models()
 
-def load_engine(model_name):
-    if model_name in _engines:
-        return _engines[model_name]
+def load_engine(model_name: str) -> InferenceEngine:
+    """Return a cached InferenceEngine for selected model"""
+    if model_name in _ENGINE_CACHE:
+        return _ENGINE_CACHE[model_name]
     model, tokenizer, config = manager.load_model(model_name)
     engine = InferenceEngine(model, tokenizer, config)
-    _engines[model_name] = engine
+    _ENGINE_CACHE[model_name] = engine
     return engine
 
 def chat_fn(message, history, model_name, max_tokens, temperature, top_p, top_k):
+    """
+    Gradio Chatbot callback.
+    - history: list of dicts [{role: "user"/"assistant", content: "..."}, ...]
+    - message: latest user message string
+    """
     if not model_name:
-        return history + [[message, "No model selected. Please choose a model."]]
+        # Append assistant message indicating the issue
+        history = history + [{"role": "assistant", "content": "No model selected. Please choose a model from the right panel."}]
+        return history
+
     try:
         engine = load_engine(model_name)
     except Exception as e:
-        return history + [[message, f"Error loading model: {e}"]]
-    reply = engine.generate_response(
-        message,
-        max_tokens=max_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        top_k=top_k
-    )
-    return history + [[message, reply]]
+        history = history + [{"role": "assistant", "content": f"Error loading model: {e}"}]
+        return history
+
+    try:
+        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"An error occurred during generation: {e}"
+
+    # Append the user and assistant messages in messages format
+    history = history + [
+        {"role": "user", "content": message},
+        {"role": "assistant", "content": reply},
+    ]
+    return history
 
 def clear_chat():
+    """Reset chat history"""
     return []
 
 with gr.Blocks(title="Automotive SLM Chatbot") as demo:
     gr.Markdown("# 🚗 Automotive SLM Chatbot (Gradio)")
+    gr.Markdown("Small Language Model for automotive assistance. Select a model and start chatting.")
+
     with gr.Row():
         with gr.Column(scale=3):
-            chatbot = gr.Chatbot(height=450, label="Chat")
-            msg = gr.Textbox(placeholder="Ask about automotive topics...", label="Your message")
+            chatbot = gr.Chatbot(
+                label="Chat",
+                height=500,
+                type="messages"  # use OpenAI-style messages
+            )
+            msg = gr.Textbox(
+                placeholder="Ask about automotive topics (e.g., tire pressure, check engine light, EV charging)...",
+                label="Your message"
+            )
             with gr.Row():
                 send_btn = gr.Button("Send", variant="primary")
                 clear_btn = gr.Button("Clear")
+
         with gr.Column(scale=2):
             gr.Markdown("### Model settings")
             available = list_models()
             if not available:
-                # Show a friendly message and stop early
-                import gradio as gr
-                with gr.Row():
-                    gr.Markdown("No models found in assets/models. Please add .pt/.pth/.onnx files and refresh.")
+                gr.Markdown("No models found in assets/models. Please add .pt/.pth/.onnx files and refresh the Space.")
+                # Disabled controls to avoid wiring errors
+                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: show size labels
+                def size_mb(path):
+                    try:
+                        return os.path.getsize(path) / (1024 * 1024)
+                    except Exception:
+                        return 0.0
+                labels = []
+                for name in available:
+                    mb = size_mb(os.path.join(MODELS_DIR, name))
+                    labels.append(f"{name} ({mb:.1f} MB)")
+                # Map labels to values so dropdown shows label but value is filename
+                choices = list(zip(labels, available))
+
                 model_dropdown = gr.Dropdown(
-                    choices=available,
+                    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")
-            gr.Markdown("Tip: lower temperature for more deterministic answers.")
-    # Events
-    send_evt = 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, inputs=None, outputs=[chatbot])
+                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")
+
+            gr.Markdown("Tip: Lower temperature and higher top-k/top-p can make answers more focused.")
+
+    # Wire events only if models are available
+    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)
 
 if __name__ == "__main__":
     demo.launch()