Update app to load models from the model repo

app.py

@@ -8,6 +8,8 @@ import time
 import asyncio
 import sys
 
+from huggingface_hub import hf_hub_download
+
 os.environ["STREAMLIT_WATCHED_FILES"] = ""
 if sys.platform.startswith("win"):
     asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
@@ -373,31 +375,33 @@ class BiLSTMLanguageModel(nn.Module):
 # -------------------------------
 @st.cache_resource
 def load_models():
-    # Load SentencePiece Model
-    sp_model_path = "urdu_sp.model"
-    sp = spm.SentencePieceProcessor()
-    
-    if not os.path.exists(sp_model_path):
-        st.error("SentencePiece model not found! Make sure urdu_sp.model is present.")
+    # Define your model repository
+    repo_id = "zaiffi/Mehfil-e-Sukhan"  # Replace with your actual username/model-name
+    
+    # Download SentencePiece Model
+    try:
+        sp_model_path = hf_hub_download(repo_id=repo_id, filename="urdu_sp.model")
+        sp = spm.SentencePieceProcessor()
+        sp.load(sp_model_path)
+    except Exception as e:
+        st.error(f"Error loading SentencePiece model: {e}")
         return None, None
     
-    sp.load(sp_model_path)
-    
     # Initialize & Load the Trained Model
     vocab_size = sp.get_piece_size()
     model = BiLSTMLanguageModel(vocab_size, embed_dim=512, hidden_dim=768, num_layers=3, dropout=0.2)
     
-    weights_path = "model_weights.pth"
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
     
-    if not os.path.exists(weights_path):
-        st.error("Trained model weights not found! Make sure model_weights.pth is present.")
+    try:
+        weights_path = hf_hub_download(repo_id=repo_id, filename="model_weights.pth")
+        model.load_state_dict(torch.load(weights_path, map_location=device))
+        model.eval()
+    except Exception as e:
+        st.error(f"Error loading model weights: {e}")
         return sp, None
     
-    model.load_state_dict(torch.load(weights_path, map_location=device))
-    model.eval()
-    
     return sp, model
 
 def generate_poetry_nucleus(model, sp, start_word, num_words=12, temperature=1.2, top_p=0.85):

requirements.txt

@@ -1,3 +1,4 @@
 torch==2.6.0
 sentencepiece==0.2.0
+huggingface-hub==0.29.3
 streamlit==1.43.0