Update app.py

app.py

@@ -1,25 +1,34 @@
 import streamlit as st
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
-from huggingface_hub import snapshot_download
+from huggingface_hub import snapshot_download, login
 from pathlib import Path
 
 def main():
     st.title("Codestral Inference with Hugging Face")
     
+    # Get the Hugging Face API token from the user
+    hf_token = st.text_input("Enter your Hugging Face API token", type="password")
+    if not hf_token:
+        st.warning("Please enter your Hugging Face API token to proceed.")
+        st.stop()
+
+    # Login to Hugging Face Hub
+    login(hf_token)
+    
     # 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)
+    snapshot_download(repo_id=model_id, allow_patterns=["*.bin", "*.json", "*.model"], local_dir=local_model_path, use_auth_token=hf_token)
     st.success("Model downloaded successfully!")
     
     # Load the model and tokenizer
     st.text("Loading model...")
-    tokenizer = AutoTokenizer.from_pretrained(local_model_path)
-    model = AutoModelForCausalLM.from_pretrained(local_model_path)
-    generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+    tokenizer = AutoTokenizer.from_pretrained(local_model_path, use_auth_token=hf_token)
+    model = AutoModelForCausalLM.from_pretrained(local_model_path, use_auth_token=hf_token)
+    generator = pipeline("text-generation", model=model, tokenizer=tokenizer, use_auth_token=hf_token)
     st.success("Model loaded successfully!")
     
     user_input = st.text_area("Enter your instruction", "Explain Machine Learning to me in a nutshell.")