Update app.py

app.py

@@ -13,16 +13,13 @@ from langchain.chat_models import ChatOpenAI
 tokenizer = AutoTokenizer.from_pretrained("HooshvareLab/bert-fa-base-uncased")
 model = AutoModel.from_pretrained("HooshvareLab/bert-fa-base-uncased")
 
+@st.cache
 def get_embedding(text):
-    sub_chunks = split_text_to_chunks(text)
-    all_embeddings = []
-    for chunk in sub_chunks:
-        inputs = tokenizer(chunk, return_tensors="pt", truncation=True, padding=True, max_length=512)
-        with torch.no_grad():
-            outputs = model(**inputs)
-        embedding = outputs.last_hidden_state.mean(dim=1).squeeze().numpy()
-        all_embeddings.append(embedding)
-    return np.mean(all_embeddings, axis=0)
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    embeddings = outputs.last_hidden_state.mean(dim=1)
+    return embeddings.squeeze().numpy()
 
 def cosine_similarity(vec1, vec2):
     return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))