Update app.py

app.py

@@ -8,7 +8,7 @@ from sklearn.metrics.pairwise import cosine_similarity
 from transformers import AutoTokenizer, AutoModel
 import torch
 from torch.utils.data import DataLoader, Dataset
-from datasets import load_dataset  # For loading dataset
+from datasets import load_dataset
 from datetime import datetime
 from typing import List, Dict, Any
 from functools import partial
@@ -24,20 +24,36 @@ if 'feedback' not in st.session_state:
     st.session_state.feedback = {}
 
 # Define subset size and batch size for optimization
-SUBSET_SIZE = 500  # Smaller subset for faster precomputation
+SUBSET_SIZE = 500  # Subset for faster precomputation
 BATCH_SIZE = 8  # Smaller batch size to reduce memory overhead
 
-# Caching key resources: Model, Tokenizer, and Precomputed Embeddings
 @st.cache_resource
-def load_model_and_tokenizer():
+def load_model_and_tokenizer_with_progress():
     """
-    Load the pre-trained model and tokenizer using Hugging Face Transformers.
-    Cached to ensure it loads only once.
+    Load the pre-trained model and tokenizer using Hugging Face Transformers
+    with a progress bar for better user experience.
     """
-    model_name = "Salesforce/codet5-small"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModel.from_pretrained(model_name).to(device)
-    model.eval()
+    progress_bar = st.progress(0)
+    status_text = st.empty()
+    
+    try:
+        progress_bar.progress(10)
+        status_text.text("Loading tokenizer...")
+        model_name = "Salesforce/codet5-small"
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        
+        progress_bar.progress(50)
+        status_text.text("Loading model...")
+        model = AutoModel.from_pretrained(model_name).to(device)
+        model.eval()
+        
+        progress_bar.progress(100)
+        status_text.text("Model loaded successfully!")
+        
+    finally:
+        progress_bar.empty()
+        status_text.empty()
+    
     return tokenizer, model
 
 @st.cache_resource
@@ -56,7 +72,6 @@ def load_data():
     data['text'] = data['docstring'].fillna('') + ' ' + data['summary'].fillna('')
     return data
 
-
 @st.cache_resource
 def precompute_embeddings(data: pd.DataFrame, _tokenizer, _model, batch_size=BATCH_SIZE):
     """
@@ -113,10 +128,13 @@ def precompute_embeddings(data: pd.DataFrame, _tokenizer, _model, batch_size=BAT
     )
 
     embeddings = []
-    for batch in dataloader:
+    progress_bar = st.progress(0)  # Progress bar for embedding computation
+    for i, batch in enumerate(dataloader):
         batch_embeddings = generate_embeddings_batch(_model, batch, device)
         embeddings.extend(batch_embeddings)
+        progress_bar.progress((i + 1) / len(dataloader))
 
+    progress_bar.empty()
     data['embedding'] = embeddings
     return data
 
@@ -136,8 +154,18 @@ def find_similar_repos(query_embedding: np.ndarray, data: pd.DataFrame, top_n=5)
     """
     Compute cosine similarity and return the top N most similar repositories.
     """
-    similarities = cosine_similarity([query_embedding], np.stack(data['embedding'].values))[0]
+    # Reshape query_embedding to 2D
+    query_embedding = query_embedding.reshape(1, -1)
+    
+    # Convert data['embedding'] to a 2D array
+    embeddings = np.vstack(data['embedding'].values)
+    
+    # Compute cosine similarity
+    similarities = cosine_similarity(query_embedding, embeddings)[0]
+    
+    # Add similarity scores to the DataFrame
     data['similarity'] = similarities
+    
     return data.nlargest(top_n, 'similarity')
 
 def display_recommendations(recommendations: pd.DataFrame):
@@ -154,8 +182,10 @@ def display_recommendations(recommendations: pd.DataFrame):
 st.title("Repository Recommender System 🚀")
 st.caption("Find repositories based on your project description.")
 
-# Load resources
-tokenizer, model = load_model_and_tokenizer()
+# Load resources with progress bar
+tokenizer, model = load_model_and_tokenizer_with_progress()
+
+# Load data and precompute embeddings
 data = load_data()
 data = precompute_embeddings(data, tokenizer, model)