Update app.py

app.py

@@ -7,11 +7,9 @@ import numpy as np
 from sklearn.metrics.pairwise import cosine_similarity
 from transformers import AutoTokenizer, AutoModel
 import torch
-from tqdm import tqdm
-from datasets import load_dataset
+from torch.utils.data import DataLoader, Dataset
 from datetime import datetime
 from typing import List, Dict, Any
-from torch.utils.data import DataLoader, Dataset
 from functools import partial
 
 # Configure GPU if available
@@ -20,325 +18,149 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 # Initialize session state
 if 'history' not in st.session_state:
     st.session_state.history = []
+
 if 'feedback' not in st.session_state:
     st.session_state.feedback = {}
 
 # Define subset size
-SUBSET_SIZE = 1000  # Starting with 500 items for quick testing
-
-class TextDataset(Dataset):
-    def __init__(self, texts: List[str], tokenizer, max_length: int = 512):
-        self.texts = texts
-        self.tokenizer = tokenizer
-        self.max_length = max_length
-
-    def __len__(self):
-        return len(self.texts)
-
-    def __getitem__(self, idx):
-        return self.tokenizer(
-            self.texts[idx],
-            padding='max_length',
-            truncation=True,
-            max_length=self.max_length,
-            return_tensors="pt"
-        )
+SUBSET_SIZE = 1000
 
-def generate_case_study(row: Dict[str, Any]) -> str:
-    """Generate a detailed case study for a repository using available metadata"""
-    # Extract relevant information from the row
-    summary = row.get('summary', '').strip()
-    docstring = row.get('docstring', '').strip()
-    repo_name = row.get('repo', '').strip()
-    
-    # Generate a more detailed overview using available information
-    overview = summary if summary else "This repository provides a software implementation"
-    if docstring:
-        # Extract the first paragraph of the docstring for additional context
-        first_para = docstring.split('\n\n')[0].strip()
-        overview = f"{overview}. {first_para}"
-    
-    # Analyze the repository path to infer technology stack
-    path_components = row.get('path', '').lower().split('/')
-    tech_stack = []
-    
-    # Common technology indicators in paths
-    if any('python' in comp for comp in path_components):
-        tech_stack.append("Python")
-    if any('tensorflow' in comp or 'tf' in comp for comp in path_components):
-        tech_stack.append("TensorFlow")
-    if any('pytorch' in comp for comp in path_components):
-        tech_stack.append("PyTorch")
-    if any('react' in comp for comp in path_components):
-        tech_stack.append("React")
-    
-    tech_stack_str = ", ".join(tech_stack) if tech_stack else "various technologies"
-    
-    case_study = f"""
-### Overview
-{overview}
-
-### Technical Implementation
-This project is built using {tech_stack_str}. The implementation focuses on providing a robust and maintainable solution for {summary.lower() if summary else 'the specified requirements'}.
-
-### Key Features
-- Primary functionality: {summary if summary else 'Implementation of core project requirements'}
-- Complete documentation and code examples
-- Well-structured implementation following best practices
-- Modular design for easy integration and customization
-
-### Use Cases
-This repository is particularly valuable for:
-- Developers implementing similar functionality in their projects
-- Teams looking for reference implementations and best practices
-- Projects requiring similar technical capabilities
-- Learning and educational purposes in related technical domains
-
-### Integration Considerations
-The repository can be integrated into existing projects, with consideration for:
-- Compatibility with existing technology stacks
-- Required dependencies and prerequisites
-- Potential customization needs
-- Performance and scalability requirements
+# Caching key resources: Model, Tokenizer, and Precomputed Embeddings
+@st.cache_resource
+def load_model_and_tokenizer():
     """
-    return case_study
-
-def display_recommendations(recommendations: pd.DataFrame):
-    """Display recommendations in a list format with all details"""
-    st.markdown("### 🎯 Top Recommendations")
-    
-    # Create a list of recommendations
-    for idx, row in recommendations.iterrows():
-        with st.container():
-            # Header with repository name and match score
-            col1, col2 = st.columns([3, 1])
-            with col1:
-                st.markdown(f"### {idx + 1}. {row['repo']}")
-            with col2:
-                st.metric("Match Score", f"{row['similarity']:.2%}")
-            
-            # Repository details
-            st.markdown(f"**URL:** [View Repository]({row['url']})")
-            st.markdown(f"**Path:** `{row['path']}`")
-            
-            # Feedback buttons
-            col1, col2, col3 = st.columns([1, 1, 4])
-            with col1:
-                if st.button("👍", key=f"like_{idx}"):
-                    st.session_state.feedback[row['repo']] = st.session_state.feedback.get(row['repo'], {'likes': 0, 'dislikes': 0})
-                    st.session_state.feedback[row['repo']]['likes'] += 1
-                    st.success("Thanks for your feedback!")
-            with col2:
-                if st.button("👎", key=f"dislike_{idx}"):
-                    st.session_state.feedback[row['repo']] = st.session_state.feedback.get(row['repo'], {'likes': 0, 'dislikes': 0})
-                    st.session_state.feedback[row['repo']]['dislikes'] += 1
-                    st.success("Thanks for your feedback!")
-            
-            # Documentation and case study in tabs
-            tab1, tab2 = st.tabs(["📚 Documentation", "📑 Case Study"])
-            with tab1:
-                if row['docstring']:
-                    st.markdown(row['docstring'])
-                else:
-                    st.info("No documentation available")
-            
-            with tab2:
-                st.markdown(generate_case_study(row))
-            
-            st.markdown("---")
+    Load the pre-trained model and tokenizer using Hugging Face Transformers.
+    Cached to ensure it loads only once.
+    """
+    model_name = "Salesforce/codet5-small"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModel.from_pretrained(model_name).to(device)
+    model.eval()
+    return tokenizer, model
 
 @st.cache_resource
-def load_data_and_model():
-    """Load the dataset and model with optimized memory usage"""
-    try:
-        # Load dataset
-        dataset = load_dataset("frankjosh/filtered_dataset")
-        data = pd.DataFrame(dataset['train'])
-        
-        # Take a random subset
-        data = data.sample(n=min(SUBSET_SIZE, len(data)), random_state=42).reset_index(drop=True)
-        
-        # Combine text fields
-        data['text'] = data['docstring'].fillna('') + ' ' + data['summary'].fillna('')
-        
-        # Load model and tokenizer
-        model_name = "Salesforce/codet5-small"
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModel.from_pretrained(model_name)
-        
-        if torch.cuda.is_available():
-            model = model.to(device)
-        
-        model.eval()
-        return data, tokenizer, model
-    
-    except Exception as e:
-        st.error(f"Error in initialization: {str(e)}")
-        st.stop()
-
-def collate_fn(batch, pad_token_id):
-    max_length = max(inputs['input_ids'].shape[1] for inputs in batch)
-    input_ids = []
-    attention_mask = []
-    
-    for inputs in batch:
-        input_ids.append(torch.nn.functional.pad(
-            inputs['input_ids'].squeeze(),
-            (0, max_length - inputs['input_ids'].shape[1]),
-            value=pad_token_id
-        ))
-        attention_mask.append(torch.nn.functional.pad(
-            inputs['attention_mask'].squeeze(),
-            (0, max_length - inputs['attention_mask'].shape[1]),
-            value=0
-        ))
-    
-    return {
-        'input_ids': torch.stack(input_ids),
-        'attention_mask': torch.stack(attention_mask)
-    }
+def load_data():
+    """
+    Load and sample the dataset from Hugging Face.
+    Returns a DataFrame with a fixed subset of repositories.
+    """
+    dataset = load_dataset("frankjosh/filtered_dataset")
+    data = pd.DataFrame(dataset['train'])
+    data = data.sample(n=min(SUBSET_SIZE, len(data)), random_state=42).reset_index(drop=True)
+    return data
 
-def generate_embeddings_batch(model, batch, device):
-    """Generate embeddings for a batch of inputs"""
-    with torch.no_grad():
-        batch = {k: v.to(device) for k, v in batch.items()}
-        outputs = model.encoder(**batch)
-        embeddings = outputs.last_hidden_state.mean(dim=1)
-        return embeddings.cpu().numpy()
+@st.cache_resource
+def precompute_embeddings(data: pd.DataFrame, tokenizer, model, batch_size=16):
+    """
+    Precompute embeddings for repository metadata to optimize query performance.
+    """
+    class TextDataset(Dataset):
+        def __init__(self, texts: List[str], tokenizer, max_length=512):
+            self.texts = texts
+            self.tokenizer = tokenizer
+            self.max_length = max_length
+
+        def __len__(self):
+            return len(self.texts)
+
+        def __getitem__(self, idx):
+            return self.tokenizer(
+                self.texts[idx],
+                padding='max_length',
+                truncation=True,
+                max_length=self.max_length,
+                return_tensors="pt"
+            )
+
+    def collate_fn(batch, pad_token_id):
+        max_length = max(inputs['input_ids'].shape[1] for inputs in batch)
+        input_ids, attention_mask = [], []
+        for inputs in batch:
+            input_ids.append(torch.nn.functional.pad(
+                inputs['input_ids'].squeeze(),
+                (0, max_length - inputs['input_ids'].shape[1]),
+                value=pad_token_id
+            ))
+            attention_mask.append(torch.nn.functional.pad(
+                inputs['attention_mask'].squeeze(),
+                (0, max_length - inputs['attention_mask'].shape[1]),
+                value=0
+            ))
+        return {
+            'input_ids': torch.stack(input_ids),
+            'attention_mask': torch.stack(attention_mask)
+        }
+
+    def generate_embeddings_batch(model, batch, device):
+        with torch.no_grad():
+            batch = {k: v.to(device) for k, v in batch.items()}
+            outputs = model.encoder(**batch)
+            return outputs.last_hidden_state.mean(dim=1).cpu().numpy()
 
-def precompute_embeddings(data: pd.DataFrame, model, tokenizer, batch_size: int = 16):
-    """Precompute embeddings with batching and progress tracking"""
     dataset = TextDataset(data['text'].tolist(), tokenizer)
     dataloader = DataLoader(
-        dataset,
-        batch_size=batch_size,
-        shuffle=False,
-        collate_fn=partial(collate_fn, pad_token_id=tokenizer.pad_token_id),
-        num_workers=2,
-        pin_memory=True
+        dataset, batch_size=batch_size, shuffle=False,
+        collate_fn=partial(collate_fn, pad_token_id=tokenizer.pad_token_id)
     )
-    
+
     embeddings = []
-    total_batches = len(dataloader)
-    
-    # Create a progress bar
-    progress_bar = st.progress(0)
-    status_text = st.empty()
-    
-    start_time = datetime.now()
-    
-    for i, batch in enumerate(dataloader):
-        # Generate embeddings for batch
+    for batch in dataloader:
         batch_embeddings = generate_embeddings_batch(model, batch, device)
         embeddings.extend(batch_embeddings)
-        
-        # Update progress
-        progress = (i + 1) / total_batches
-        progress_bar.progress(progress)
-        
-        # Calculate and display ETA
-        elapsed_time = (datetime.now() - start_time).total_seconds()
-        eta = (elapsed_time / (i + 1)) * (total_batches - (i + 1))
-        status_text.text(f"Processing batch {i+1}/{total_batches}. ETA: {int(eta)} seconds")
-    
-    progress_bar.empty()
-    status_text.empty()
-    
-    # Add embeddings to dataframe
+
     data['embedding'] = embeddings
     return data
 
 @torch.no_grad()
 def generate_query_embedding(model, tokenizer, query: str) -> np.ndarray:
-    """Generate embedding for a single query"""
+    """
+    Generate embedding for a user query using the pre-trained model.
+    """
     inputs = tokenizer(
-        query,
-        return_tensors="pt",
-        padding=True,
-        truncation=True,
-        max_length=512
+        query, return_tensors="pt", padding=True,
+        truncation=True, max_length=512
     ).to(device)
-    
     outputs = model.encoder(**inputs)
-    embedding = outputs.last_hidden_state.mean(dim=1).cpu().numpy()
-    return embedding.squeeze()
+    return outputs.last_hidden_state.mean(dim=1).cpu().numpy()
 
-def find_similar_repos(query_embedding: np.ndarray, data: pd.DataFrame, top_n: int = 5) -> pd.DataFrame:
-    """Find similar repositories using vectorized operations"""
+def find_similar_repos(query_embedding: np.ndarray, data: pd.DataFrame, top_n=5) -> pd.DataFrame:
+    """
+    Compute cosine similarity and return the top N most similar repositories.
+    """
     similarities = cosine_similarity([query_embedding], np.stack(data['embedding'].values))[0]
     data['similarity'] = similarities
     return data.nlargest(top_n, 'similarity')
 
-# Load resources
-data, tokenizer, model = load_data_and_model()
-
-# Add info about subset size
-st.info(f"Running with a subset of {SUBSET_SIZE} repositories for testing purposes.")
-
-# Precompute embeddings for the subset
-data = precompute_embeddings(data, model, tokenizer)
+def display_recommendations(recommendations: pd.DataFrame):
+    """
+    Display the recommended repositories in the Streamlit app interface.
+    """
+    st.markdown("### 🎯 Top Recommendations")
+    for idx, row in recommendations.iterrows():
+        st.markdown(f"### {idx + 1}. {row['repo']}")
+        st.metric("Match Score", f"{row['similarity']:.2%}")
+        st.markdown(f"[View Repository]({row['url']})")
 
-# Main App Interface
+# Main workflow
 st.title("Repository Recommender System 🚀")
-st.caption("Testing Version - Running on subset of data")
+st.caption("Find repositories based on your project description.")
 
-# Main interface
+# Load resources
+tokenizer, model = load_model_and_tokenizer()
+data = load_data()
+data = precompute_embeddings(data, tokenizer, model)
+
+# User input
 user_query = st.text_area(
-    "Describe your project:",
-    height=150,
-    placeholder="Example: I need a machine learning project for customer churn prediction..."
+    "Describe your project:", height=150,
+    placeholder="Example: A machine learning project for customer churn prediction..."
 )
 
-# Search button and filters
-col1, col2 = st.columns([2, 1])
-with col1:
-    search_button = st.button("🔍 Search Repositories", type="primary")
-with col2:
-    top_n = st.selectbox("Number of results:", [3, 5, 10], index=1)
-
-if search_button and user_query.strip():
-    with st.spinner("Finding relevant repositories..."):
-        # Generate query embedding and get recommendations
-        query_embedding = generate_query_embedding(model, tokenizer, user_query)
-        recommendations = find_similar_repos(query_embedding, data, top_n)
-        
-        # Save to history
-        st.session_state.history.append({
-            'query': user_query,
-            'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-            'results': recommendations['repo'].tolist()
-        })
-        
-        # Display recommendations using the new function
-        display_recommendations(recommendations)
-
-# Sidebar for History and Stats
-with st.sidebar:
-    st.header("📊 Search History")
-    if st.session_state.history:
-        for idx, item in enumerate(reversed(st.session_state.history[-5:])):
-            st.markdown(f"**Search {len(st.session_state.history)-idx}**")
-            st.markdown(f"Query: _{item['query'][:30]}..._")
-            st.caption(f"Time: {item['timestamp']}")
-            st.caption(f"Results: {len(item['results'])} repositories")
-            if st.button("Rerun this search", key=f"rerun_{idx}"):
-                st.session_state.rerun_query = item['query']
-            st.markdown("---")
+if st.button("🔍 Search Repositories"):
+    if user_query.strip():
+        with st.spinner("Finding relevant repositories..."):
+            query_embedding = generate_query_embedding(model, tokenizer, user_query)
+            recommendations = find_similar_repos(query_embedding, data)
+            display_recommendations(recommendations)
     else:
-        st.write("No search history yet")
-
-    st.header("📈 Usage Statistics")
-    st.write(f"Total Searches: {len(st.session_state.history)}")
-    if st.session_state.feedback:
-        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']])
-
-# Footer
-st.markdown("---")
-st.markdown(
-    """
-    Made with 🤖 using CodeT5 and Streamlit |
-    
-    """
-)
+        st.error("Please provide a project description.")