Update app.py

app.py

@@ -1,13 +1,24 @@
 import streamlit as st
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from huggingface_hub import snapshot_download
+from pathlib import Path
 
 def main():
     st.title("Codestral Inference with Hugging Face")
     
+    # Download the model files
+    st.text("Downloading model...")
+    model_id = "mistralai/Codestral-22B-v0.1"
+    local_model_path = Path.home().joinpath('mistral_models', model_id)
+    local_model_path.mkdir(parents=True, exist_ok=True)
+
+    snapshot_download(repo_id=model_id, allow_patterns=["*.bin", "*.json", "*.model"], local_dir=local_model_path)
+    st.success("Model downloaded successfully!")
+    
     # Load the model and tokenizer
     st.text("Loading model...")
-    tokenizer = AutoTokenizer.from_pretrained("mistralai/Codestral-22B-v0.1")
-    model = AutoModelForCausalLM.from_pretrained("mistralai/Codestral-22B-v0.1")
+    tokenizer = AutoTokenizer.from_pretrained(local_model_path)
+    model = AutoModelForCausalLM.from_pretrained(local_model_path)
     generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
     st.success("Model loaded successfully!")