Update app.py

app.py

@@ -1,20 +1,21 @@
 import gradio as gr
 from huggingface_hub import snapshot_download
 from pathlib import Path
+import spaces
 from mistral.cli.chat import load_model, generate_stream
 
-# Download the model
 mistral_models_path = Path.home().joinpath('mistral_models', 'mamba-codestral-7B-v0.1')
 mistral_models_path.mkdir(parents=True, exist_ok=True)
 
 snapshot_download(repo_id="mistralai/mamba-codestral-7B-v0.1", 
                   allow_patterns=["params.json", "consolidated.safetensors", "tokenizer.model.v3"], 
                   local_dir=mistral_models_path)
+MODEL_PATH = str(mistral_models_path)
 
-# Load the model
-model = load_model(str(mistral_models_path))
 
+@spaces.GPU()
 def generate_response(message, history):
+    model = load_model(MODEL_PATH)
     history_mistral_format = [
         {"role": "user" if i % 2 == 0 else "assistant", "content": m}
         for i, m in enumerate(sum(history, []))
@@ -24,12 +25,17 @@ def generate_response(message, history):
     response = ""
     for chunk in generate_stream(model, history_mistral_format, max_tokens=256):
         response += chunk
-        yield response
+    return response
+
+# Gradio interface
+def chat_interface(message, history):
+    response = generate_response(message, history, model)
+    return response
 
 iface = gr.ChatInterface(
-    generate_response,
-    title="Mamba Codestral Chat",
-    description="Chat with the Mamba Codestral 7B model.",
+    chat_interface,
+    title="Mamba Codestral Chat (ZeroGPU)",
+    description="Chat with the Mamba Codestral 7B model using Hugging Face Spaces ZeroGPU feature.",
 )
 
 if __name__ == "__main__":