Update src/streamlit_app.py

src/streamlit_app.py

@@ -4,6 +4,7 @@ import torch
 import pickle
 import numpy as np
 from transformers import AutoTokenizer, AutoModel
+from sentence_transformers import SentenceTransformer
 from sklearn.feature_extraction.text import ENGLISH_STOP_WORDS
 from PyPDF2 import PdfReader
 import json
@@ -52,7 +53,7 @@ JSON_TEMPLATE = {
 
 # --- Load Model ---
 tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
-model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
+model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
 
 def mean_pooling(model_output, attention_mask):
     token_embeddings = model_output[0]
@@ -60,26 +61,24 @@ def mean_pooling(model_output, attention_mask):
     return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
 
 def embed_text(texts):
-    # encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors='pt')
-    # with torch.no_grad():
-    #     model_output = model(**encoded_input)
-    # return mean_pooling(model_output, encoded_input['attention_mask']).cpu().numpy()
     
-    # Use the correct device
-    device = torch.device("cpu")  # or "cuda" if you have a GPU
+    # # Use the correct device
+    # device = torch.device("cpu")  # or "cuda" if you have a GPU
 
-    # Move model to the correct device
-    model.to(device)
+    # # Move model to the correct device
+    # model.to(device)
 
-    # Prepare the inputs
-    encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors='pt').to(device)
+    # # Prepare the inputs
+    # encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors='pt').to(device)
 
-    # Generate embeddings
-    with torch.no_grad():
-        model_output = model(**encoded_input)
+    # # Generate embeddings
+    # with torch.no_grad():
+    #     model_output = model(**encoded_input)
 
-    # Return embeddings as numpy array
-    return mean_pooling(model_output, encoded_input['attention_mask']).cpu().numpy()
+    # # Return embeddings as numpy array
+    # return mean_pooling(model_output, encoded_input['attention_mask']).cpu().numpy()
+    embeddings = model.encode(texts)
+    return embeddings
 
 
 # --- Load & Chunk Text ---