Update app.py

app.py

@@ -35,14 +35,26 @@ if 'history' not in st.session_state:
 if 'feedback' not in st.session_state:
     st.session_state.feedback = {}
 
-
-
 @st.cache_data
-def generate_embedding(_model, tokenizer, text):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
+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()}
     with torch.no_grad():
-        outputs = model(**inputs)
-    return outputs.last_hidden_state.mean(dim=1).squeeze().numpy()
+        outputs = _model.encoder(**inputs)
+    embedding = outputs.last_hidden_state.mean(dim=1).squeeze()
+    if torch.cuda.is_available():
+        embedding = embedding.cpu()
+    return embedding.numpy()
+
+#error handling
+try:
+    query_embedding = generate_embedding(model, tokenizer, user_query)
+except Exception as e:
+    st.error(f"Error generating embedding: {str(e)}")
+    st.stop()
+    
+   
 
 # Step 1: Load Dataset and Precompute Embeddings
 @st.cache_resource