Update app.py

app.py

@@ -1,8 +1,3 @@
-# -*- coding: utf-8 -*-
-"""app.py
-Enhanced Repository Recommender System using Streamlit and CodeT5-small
-"""
-
 import warnings
 warnings.filterwarnings('ignore')
 
@@ -15,6 +10,9 @@ import torch
 from tqdm import tqdm
 from datasets import load_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
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -25,94 +23,163 @@ 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
+# Define subset size
+SUBSET_SIZE = 1000  # Starting with 1000 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"
+        )
+
 @st.cache_resource
 def load_data_and_model():
-    """
-    Load the dataset and precompute embeddings. Load the CodeT5-small model and tokenizer.
-    """
+    """Load the dataset and model with optimized memory usage"""
     try:
-        # Download and load dataset
+        # 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
+        
+        # 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('')
-    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"
-    try:
+        
+        # Load model and tokenizer
+        model_name = "Salesforce/codet5-small"
         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
+            model = model.to(device)
+        
+        model.eval()
+        return data, tokenizer, model
+    
     except Exception as e:
-        st.error(f"Error loading model: {str(e)}")
+        st.error(f"Error in initialization: {str(e)}")
         st.stop()
 
-    return data, tokenizer, model
+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)
+    }
 
-# Define the embedding generation function
-@st.cache_data
-def generate_embedding(_model, _tokenizer, text):
-    inputs = _tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
-    if torch.cuda.is_available():
-        inputs = {k: v.to('cuda') for k, v in inputs.items()}
+def generate_embeddings_batch(model, batch, device):
+    """Generate embeddings for a batch of inputs"""
     with torch.no_grad():
-        outputs = _model.encoder(**inputs)
-    embedding = outputs.last_hidden_state.mean(dim=1).squeeze()
-    if torch.cuda.is_available():
-        embedding = embedding.cpu()
-    return embedding.numpy()
-
-# Precompute embeddings for dataset
-def precompute_embeddings(data, model, tokenizer):
+        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()
+
+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,  # Reduced workers for smaller dataset
+        pin_memory=True
+    )
+    
     embeddings = []
-    for text in tqdm(data['text'], desc="Generating embeddings"):
-        embedding = generate_embedding(model, tokenizer, text)
-        embeddings.append(embedding)
+    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
+        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
 
-# Generate a concise case study brief from repository data
-def generate_case_study(repo_data):
-    template = f"""
-    **Project Overview**: {repo_data['summary'][:50]}...
+@torch.no_grad()
+def generate_query_embedding(model, tokenizer, query: str) -> np.ndarray:
+    """Generate embedding for a single query"""
+    inputs = tokenizer(
+        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()
 
-    **Key Features**:
-    - Repository contains production-ready {repo_data['path'].split('/')[-1]} implementation
-    - {repo_data['docstring'][:50]}...
+def find_similar_repos(query_embedding: np.ndarray, data: pd.DataFrame, top_n: int = 5) -> pd.DataFrame:
+    """Find similar repositories using vectorized operations"""
+    similarities = cosine_similarity([query_embedding], np.stack(data['embedding'].values))[0]
+    data['similarity'] = similarities
+    return data.nlargest(top_n, 'similarity')
 
-    **Potential Applications**: This repository can be utilized for projects requiring {' '.join(repo_data['summary'].split()[:5])}...
+# Load resources
+data, tokenizer, model = load_data_and_model()
 
-    **Implementation Complexity**: {'Medium' if len(repo_data['docstring']) > 500 else 'Low'}
+# Add info about subset size
+st.info(f"Running with a subset of {SUBSET_SIZE} repositories for testing purposes.")
 
-    **Integration Potential**: {'High' if 'api' in repo_data['text'].lower() or 'interface' in repo_data['text'].lower() else 'Medium'}
-    """
-    return template[:150] + "..."
+# Precompute embeddings for the subset
+data = precompute_embeddings(data, model, tokenizer)
 
-# Save user feedback for a repository
-def save_feedback(repo_id, feedback_type):
-    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 Interface
+st.title("Repository Recommender System 🚀")
+st.caption("Testing Version - Running on subset of data")
 
-# Load resources
-data, tokenizer, model = load_data_and_model()
-data = precompute_embeddings(data, model, tokenizer)
+# Rest of your UI code remains the same...
 
 # Main App Interface
 st.title("Enhanced Repository Recommender System 🚀")