Update app.py

app.py

@@ -78,6 +78,8 @@ def load_data_and_model():
     
     tokenizer, model = load_model_and_tokenizer()
 
+
+# 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)
@@ -90,34 +92,26 @@ def generate_embedding(_model, _tokenizer, text):
         embedding = embedding.cpu()
     return embedding.numpy()
 
-#error handling
+# 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 with GPU support
+# Precompute embeddings for dataset
+def precompute_embeddings(data, model, tokenizer):
     @st.cache_data
-    def generate_embedding(text):
-        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
-        # Move inputs to GPU if available
-        if torch.cuda.is_available():
-            inputs = {k: v.to('cuda') for k, v in inputs.items()}
-        with torch.no_grad():
-            outputs = model.encoder(**inputs)
-        # Move output back to CPU if needed
-        embedding = outputs.last_hidden_state.mean(dim=1).squeeze()
-        if torch.cuda.is_available():
-            embedding = embedding.cpu()
-        return embedding.numpy()
-
-    # Generate embeddings with progress bar
+    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_embedding(x))
-    
-return data, tokenizer, model
+        data['embedding'] = data['text'].apply(lambda x: generate_cached_embedding(x))
+    return data
 
+# Example usage:
+# data = precompute_embeddings(data, model, tokenizer)
 
 def generate_case_study(repo_data):
     """