Update app.py

app.py

@@ -3,12 +3,12 @@ warnings.filterwarnings('ignore')
 
 import streamlit as st
 import pandas as pd
-from datasets import load_dataset
 import numpy as np
 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 datetime import datetime
 from typing import List, Dict, Any
 from functools import partial
@@ -23,8 +23,9 @@ if 'history' not in st.session_state:
 if 'feedback' not in st.session_state:
     st.session_state.feedback = {}
 
-# Define subset size
-SUBSET_SIZE = 1000
+# Define subset size and batch size for optimization
+SUBSET_SIZE = 500  # Smaller 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
@@ -51,9 +52,10 @@ def load_data():
     return data
 
 @st.cache_resource
-def precompute_embeddings(data: pd.DataFrame, tokenizer, model, batch_size=16):
+def precompute_embeddings(data: pd.DataFrame, _tokenizer, model, batch_size=BATCH_SIZE):
     """
     Precompute embeddings for repository metadata to optimize query performance.
+    The tokenizer is excluded from caching as it is unhashable.
     """
     class TextDataset(Dataset):
         def __init__(self, texts: List[str], tokenizer, max_length=512):
@@ -98,10 +100,10 @@ def precompute_embeddings(data: pd.DataFrame, tokenizer, model, batch_size=16):
             outputs = model.encoder(**batch)
             return outputs.last_hidden_state.mean(dim=1).cpu().numpy()
 
-    dataset = TextDataset(data['text'].tolist(), tokenizer)
+    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)
+        collate_fn=partial(collate_fn, pad_token_id=_tokenizer.pad_token_id)
     )
 
     embeddings = []