Update app.py

app.py

@@ -3,6 +3,7 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import os
 from threading import Thread
+import requests
 
 # Define model path for caching (Avoids reloading every app restart)
 MODEL_PATH = "/mnt/data/Phi-4-Hindi"
@@ -13,21 +14,28 @@ MODEL_NAME = "large-traversaal/Phi-4-Hindi"
 @st.cache_resource()
 def load_model():
     with st.spinner("Loading model... Please wait ⏳"):
-        if not os.path.exists(MODEL_PATH):
-            model = AutoModelForCausalLM.from_pretrained(
-                MODEL_NAME, token=TOKEN, trust_remote_code=True, torch_dtype=torch.bfloat16
-            )
-            tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=TOKEN)
-            model.save_pretrained(MODEL_PATH)
-            tokenizer.save_pretrained(MODEL_PATH)
-        else:
-            model = AutoModelForCausalLM.from_pretrained(MODEL_PATH)
-            tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+        try:
+            if not os.path.exists(MODEL_PATH):
+                model = AutoModelForCausalLM.from_pretrained(
+                    MODEL_NAME, token=TOKEN, trust_remote_code=True, torch_dtype=torch.bfloat16
+                )
+                tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=TOKEN)
+                model.save_pretrained(MODEL_PATH)
+                tokenizer.save_pretrained(MODEL_PATH)
+            else:
+                model = AutoModelForCausalLM.from_pretrained(MODEL_PATH)
+                tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+        except requests.exceptions.ConnectionError:
+            st.error("⚠️ Connection error! Unable to download the model. Please check your internet connection and try again.")
+            return None, None
     
     return model, tokenizer
 
 # Load and move model to appropriate device
 model, tok = load_model()
+if model is None or tok is None:
+    st.stop()
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 try:
     model = model.to(device)
@@ -117,7 +125,7 @@ if st.button("Send"):
             for output in response_generator:
                 final_response = output  # Store latest output
         
-        #st.success("✅ Response generated!")
+        st.success("✅ Response generated!")
         # Add generated response to session state
         st.experimental_rerun()