Update app.py

app.py

@@ -1,17 +1,57 @@
 import gradio as gr
 from pathlib import Path
+import os
+from huggingface_hub import snapshot_download
 from mistral_inference.transformer import Transformer
 from mistral_inference.generate import generate
 from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
 from mistral_common.protocol.instruct.messages import UserMessage, AssistantMessage, SystemMessage
 from mistral_common.protocol.instruct.request import ChatCompletionRequest
 
+def download_mistral_model():
+    """Download Mistral model if not already present."""
+    print("Checking for Mistral model...")
+    mistral_models_path = Path.home().joinpath('mistral_models', 'Nemo-Instruct')
+    
+    # Check if model files already exist
+    required_files = ["params.json", "consolidated.safetensors", "tekken.json"]
+    files_exist = all(
+        mistral_models_path.joinpath(file).exists() 
+        for file in required_files
+    )
+    
+    if not files_exist:
+        print("Downloading Mistral model (this may take a while)...")
+        mistral_models_path.mkdir(parents=True, exist_ok=True)
+        
+        snapshot_download(
+            repo_id="mistralai/Mistral-Nemo-Instruct-2407",
+            allow_patterns=required_files,
+            local_dir=mistral_models_path
+        )
+        print("Model downloaded successfully!")
+    else:
+        print("Mistral model already downloaded.")
+    
+    return mistral_models_path
+
 def setup_mistral():
     """Initialize Mistral model and tokenizer."""
-    mistral_models_path = Path.home().joinpath('mistral_models', 'Nemo-Instruct')
+    mistral_models_path = download_mistral_model()
+    print("Initializing Mistral model and tokenizer...")
     tokenizer = MistralTokenizer.from_file(f"{mistral_models_path}/tekken.json")
     model = Transformer.from_folder(mistral_models_path)
     return model, tokenizer
+
+# Global variables for model and tokenizer
+global_model = None
+global_tokenizer = None
+
+def initialize_globals():
+    """Initialize global model and tokenizer if not already done."""
+    global global_model, global_tokenizer
+    if global_model is None or global_tokenizer is None:
+        global_model, global_tokenizer = setup_mistral()
     
 def check_custom_responses(message: str) -> str:
     """Check for specific patterns and return custom responses."""
@@ -162,24 +202,6 @@ def is_image_request(message: str) -> bool:
     message_lower = message.lower()
     return any(trigger in message_lower for trigger in image_triggers)
 
-def generate_image(prompt: str) -> str:
-    """Generate an image using DALLE-4K model."""
-    try:
-        response = image_client.text_to_image(
-            prompt,
-            parameters={
-                "negative_prompt": "(worst quality, low quality, illustration, 3d, 2d, painting, cartoons, sketch), open mouth",
-                "num_inference_steps": 30,
-                "guidance_scale": 7.5,
-                "sampling_steps": 15,
-                "upscaler": "4x-UltraSharp",
-                "denoising_strength": 0.5,
-            }
-        )
-        return response
-    except Exception as e:
-        print(f"Image generation error: {e}")
-        return None
 def create_mistral_messages(history, system_message, current_message):
     """Convert chat history to Mistral message format."""
     messages = []
@@ -199,7 +221,7 @@ def create_mistral_messages(history, system_message, current_message):
     messages.append(UserMessage(content=current_message))
     
     return messages
-    
+
 def respond(message, history, system_message, max_tokens=16343, temperature=0.7, top_p=0.95):
     """Main response function using Mistral model."""
     # First check for custom responses
@@ -214,8 +236,8 @@ def respond(message, history, system_message, max_tokens=16343, temperature=0.7,
         return
 
     try:
-        # Get or initialize Mistral model and tokenizer
-        model, tokenizer = setup_mistral()
+        # Initialize global model and tokenizer if needed
+        initialize_globals()
         
         # Prepare messages for Mistral
         mistral_messages = create_mistral_messages(history, system_message, message)
@@ -224,20 +246,20 @@ def respond(message, history, system_message, max_tokens=16343, temperature=0.7,
         completion_request = ChatCompletionRequest(messages=mistral_messages)
         
         # Encode the request
-        tokens = tokenizer.encode_chat_completion(completion_request).tokens
+        tokens = global_tokenizer.encode_chat_completion(completion_request).tokens
         
         # Generate response
         out_tokens, _ = generate(
             [tokens],
-            model,
+            global_model,
             max_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,
-            eos_id=tokenizer.instruct_tokenizer.tokenizer.eos_id
+            eos_id=global_tokenizer.instruct_tokenizer.tokenizer.eos_id
         )
         
         # Decode and yield response
-        response = tokenizer.decode(out_tokens[0])
+        response = global_tokenizer.decode(out_tokens[0])
         yield response
 
     except Exception as e:
@@ -254,38 +276,48 @@ body, .gradio-container {
 # System message
 system_message = """Xylaria (v1.2.9) is an AI assistant developed by Sk Md Saad Amin, designed to provide efficient, practical support in various domains with adaptable communication."""
 
-# Create Gradio interface
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(
-            value=system_message,
-            visible=False,
-        ),
-        gr.Slider(
-            minimum=1,
-            maximum=16343,
-            value=16343,
-            step=1,
-            label="Max new tokens"
-        ),
-        gr.Slider(
-            minimum=0.1,
-            maximum=4.0,
-            value=0.7,
-            step=0.1,
-            label="Temperature"
-        ),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)"
-        ),
-    ],
-    css=custom_css
-)
+def main():
+    print("Starting Mistral Chat Interface...")
+    print("Initializing model (this may take a few minutes on first run)...")
+    
+    # Initialize model and tokenizer at startup
+    initialize_globals()
+    
+    # Create Gradio interface
+    demo = gr.ChatInterface(
+        respond,
+        additional_inputs=[
+            gr.Textbox(
+                value=system_message,
+                visible=False,
+            ),
+            gr.Slider(
+                minimum=1,
+                maximum=16343,
+                value=16343,
+                step=1,
+                label="Max new tokens"
+            ),
+            gr.Slider(
+                minimum=0.1,
+                maximum=4.0,
+                value=0.7,
+                step=0.1,
+                label="Temperature"
+            ),
+            gr.Slider(
+                minimum=0.1,
+                maximum=1.0,
+                value=0.95,
+                step=0.05,
+                label="Top-p (nucleus sampling)"
+            ),
+        ],
+        css=custom_css
+    )
+    
+    print("Launch successful! Interface is ready to use.")
+    demo.launch()
 
 if __name__ == "__main__":
-    demo.launch()
+    main()