Update app.py

app.py

@@ -6,7 +6,7 @@ import requests
 import numpy as np
 import torch
 from sklearn.neighbors import NearestNeighbors
-from transformers import AutoTokenizer, AutoModel
+from sentence_transformers import SentenceTransformer
 
 # --- CONFIGURATION ---
 HF_TOKEN = os.getenv("HF_TOKEN", "").strip()
@@ -47,11 +47,8 @@ def chunk_text(text, max_chunk_length=500):
     return chunks
 
 def embed_texts(texts):
-    encoded = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
-    with torch.no_grad():
-        model_output = model(**encoded)
-    embeddings = model_output.last_hidden_state.mean(dim=1)
-    return embeddings.cpu().numpy()
+    return model.encode(texts, convert_to_numpy=True, normalize_embeddings=True)
+
 
 def save_cache(embeddings, chunks):
     np.save(EMBEDDING_CACHE_FILE, embeddings)
@@ -106,7 +103,9 @@ def respond(message, history):
         response.raise_for_status()
         output = response.json()
         generated_text = output[0]["generated_text"]
-        answer = generated_text.split("Answer:")[-1].strip()
+        match = re.search(r"Answer:(.*)", generated_text, re.DOTALL)
+        answer = match.group(1).strip() if match else generated_text.strip()
+
 
     except Exception as e:
         print("API Error:", e)
@@ -121,8 +120,7 @@ def respond(message, history):
 # --- INIT SECTION ---
 
 # Load tokenizer and model for embeddings
-tokenizer = AutoTokenizer.from_pretrained(EMBEDDING_MODEL)
-model = AutoModel.from_pretrained(EMBEDDING_MODEL)
+model = SentenceTransformer(EMBEDDING_MODEL)
 
 # Try to load cached embeddings and chunks
 chunk_embeddings, chunks = load_cache()