Update app.py

app.py

@@ -1,14 +1,8 @@
 # -*- coding: utf-8 -*-
-"""app.ipynb
-
-Automatically generated by Colab.
-
-Original file is located at
-    https://colab.research.google.com/drive/1deINvEblsMkv9h0gJzuGB4uSamW0DMX5
+"""app.py
+Enhanced Repository Recommender System using Streamlit and CodeT5-small
 """
 
-#pip install streamlit transformers gdown torch pandas numpy
-
 import warnings
 warnings.filterwarnings('ignore')
 
@@ -18,13 +12,9 @@ import numpy as np
 from sklearn.metrics.pairwise import cosine_similarity
 from transformers import AutoTokenizer, AutoModel
 import torch
-import gdown
-from pathlib import Path
-from datetime import datetime
-import json
-import torch.cuda
-import os
+from tqdm import tqdm
 from datasets import load_dataset
+from datetime import datetime
 
 # Configure GPU if available
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -35,10 +25,6 @@ if 'history' not in st.session_state:
 if 'feedback' not in st.session_state:
     st.session_state.feedback = {}
 
-
-    
-   
-
 # Step 1: Load Dataset and Precompute Embeddings
 @st.cache_resource
 def load_data_and_model():
@@ -49,35 +35,35 @@ def load_data_and_model():
         # Download and load dataset
         dataset = load_dataset("frankjosh/filtered_dataset")
         data = pd.DataFrame(dataset['train'])
+
+        # Ensure required columns exist
+        required_columns = ['docstring', 'summary']
+        for col in required_columns:
+            if col not in data.columns:
+                st.error(f"Missing required column: {col}")
+                st.stop()
+
+        # Combine text fields for embedding generation
+        data['text'] = data['docstring'].fillna('') + ' ' + data['summary'].fillna('')
     except Exception as e:
         st.error(f"Error loading dataset: {str(e)}")
         st.stop()
 
     # Load CodeT5-small model and tokenizer
     model_name = "Salesforce/codet5-small"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoTokenizer.from_pretrained(model_name)
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModel.from_pretrained(model_name)
 
-    # Combine text fields for embedding generation
-    data['text'] = data['docstring'].fillna('') + ' ' + data['summary'].fillna('')
-    return data, tokenizer, model
-    
-    @st.cache_resource
-    def load_model_and_tokenizer():
-        try:
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
-            model = AutoModel.from_pretrained(model_name)
-            # Move model to GPU if available
-            if torch.cuda.is_available():
-                model = model.to('cuda')
-            model.eval()  # Set to evaluation mode
-            return tokenizer, model
-        except Exception as e:
-            st.error(f"Error loading model: {str(e)}")
-            st.stop()
-    
-    tokenizer, model = load_model_and_tokenizer()
+        # Move model to GPU if available
+        if torch.cuda.is_available():
+            model = model.to('cuda')
+        model.eval()  # Set to evaluation mode
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        st.stop()
 
+    return data, tokenizer, model
 
 # Define the embedding generation function
 @st.cache_data
@@ -92,31 +78,17 @@ def generate_embedding(_model, _tokenizer, text):
         embedding = embedding.cpu()
     return embedding.numpy()
 
-# Error handling for generating query embeddings
-try:
-    query_embedding = generate_embedding(model, tokenizer, user_query)
-except Exception as e:
-    st.error(f"Error generating embedding: {str(e)}")
-    st.stop()
-
 # Precompute embeddings for dataset
 def precompute_embeddings(data, model, tokenizer):
-    @st.cache_data
-    def generate_cached_embedding(text):
-        return generate_embedding(model, tokenizer, text)
-
-    # Apply embedding generation with progress bar
-    with st.spinner('Generating embeddings... This might take a few minutes on first run...'):
-        data['embedding'] = data['text'].apply(lambda x: generate_cached_embedding(x))
+    embeddings = []
+    for text in tqdm(data['text'], desc="Generating embeddings"):
+        embedding = generate_embedding(model, tokenizer, text)
+        embeddings.append(embedding)
+    data['embedding'] = embeddings
     return data
 
-# Example usage:
-# data = precompute_embeddings(data, model, tokenizer)
-
+# Generate a concise case study brief from repository data
 def generate_case_study(repo_data):
-    """
-    Generate a concise case study brief from repository data
-    """
     template = f"""
     **Project Overview**: {repo_data['summary'][:50]}...
 
@@ -124,7 +96,7 @@ def generate_case_study(repo_data):
     - Repository contains production-ready {repo_data['path'].split('/')[-1]} implementation
     - {repo_data['docstring'][:50]}...
 
-    **Potential Applications**: This repository can be utilized for projects requiring {repo_data['summary'].split()[0:5]}...
+    **Potential Applications**: This repository can be utilized for projects requiring {' '.join(repo_data['summary'].split()[:5])}...
 
     **Implementation Complexity**: {'Medium' if len(repo_data['docstring']) > 500 else 'Low'}
 
@@ -132,15 +104,17 @@ def generate_case_study(repo_data):
     """
     return template[:150] + "..."
 
+# Save user feedback for a repository
 def save_feedback(repo_id, feedback_type):
-    """
-    Save user feedback for a repository
-    """
     if repo_id not in st.session_state.feedback:
         st.session_state.feedback[repo_id] = {'likes': 0, 'dislikes': 0}
     st.session_state.feedback[repo_id][feedback_type] += 1
 
-# Main App
+# Load resources
+data, tokenizer, model = load_data_and_model()
+data = precompute_embeddings(data, model, tokenizer)
+
+# Main App Interface
 st.title("Enhanced Repository Recommender System 🚀")
 
 # Sidebar for History and Stats
@@ -159,18 +133,9 @@ with st.sidebar:
     st.header("📈 Usage Statistics")
     st.write(f"Total Searches: {len(st.session_state.history)}")
     if st.session_state.feedback:
-        total_likes = sum(f['likes'] for f in st.session_state.feedback.values())
-        total_dislikes = sum(f['dislikes'] for f in st.session_state.feedback.values())
-        st.write(f"Total Likes: {total_likes}")
-        st.write(f"Total Dislikes: {total_dislikes}")
-
-# Load resources
-@st.cache_resource
-def initialize_resources():
-    data, tokenizer, model = load_data_and_model()
-    return data, tokenizer, model
-
-data, tokenizer, model = initialize_resources()
+        feedback_df = pd.DataFrame(st.session_state.feedback).T
+        feedback_df['Total'] = feedback_df['likes'] + feedback_df['dislikes']
+        st.bar_chart(feedback_df[['likes', 'dislikes']])
 
 # Main interface
 user_query = st.text_area(
@@ -186,7 +151,7 @@ with col1:
 with col2:
     top_n = st.selectbox("Number of results:", [3, 5, 10], index=1)
 
-if search_button and user_query:
+if search_button and user_query.strip():
     with st.spinner("Finding relevant repositories..."):
         # Generate query embedding and get recommendations
         query_embedding = generate_embedding(model, tokenizer, user_query)
@@ -242,4 +207,4 @@ st.markdown(
     GPU Status: {'🟢 Enabled' if torch.cuda.is_available() else '🔴 Disabled'} |
     Model: CodeT5-Small
     """
-)
+)