Update app.py

app.py

@@ -64,36 +64,31 @@ class GroqLLM(OpenAI):
 llm = GroqLLM(api_key=groq_api_key, model_name="deepseek-r1-distill-llama-70b")
 
 # ----------------- ساخت SimpleRetriever -----------------
-class SimpleRetriever:
+class SimpleRetriever(BaseRetriever):
     def __init__(self, documents, embeddings):
         self.documents = documents
         self.embeddings = embeddings
 
-    def retrieve(self, query, top_k=1):
-        # تبدیل سوال به امبدینگ با استفاده از FarsiBERT
-        inputs = tokenizer(query, return_tensors="pt", padding=True, truncation=True)
-        with torch.no_grad():
-            query_embedding = model(**inputs).last_hidden_state.mean(dim=1).numpy()
-
+    def _get_relevant_documents(self, query):
+        query_embedding = nlp(query).vector  # تبدیل سوال به امبدینگ با استفاده از spaCy
         similarities = []
         for doc_embedding in self.embeddings:
-            similarity = query_embedding.dot(doc_embedding)  # محاسبه شباهت
+            similarity = query_embedding.dot(doc_embedding)  # محاسبه شباهت بین امبدینگ‌ها
             similarities.append(similarity)
-
-        # یافتن مستندات مشابه
+        
+        # یافتن مستندات مشابه بر اساس بیشترین شباهت
         ranked_docs = sorted(zip(similarities, self.documents), reverse=True)
-        return [doc for _, doc in ranked_docs[:top_k]]
+        return [doc for _, doc in ranked_docs[:5]]  # بازگرداندن 5 مستند مشابه
+
 
-# استفاده از retriever برای جستجو
 retriever = SimpleRetriever(documents, embeddings)
 
 chain = RetrievalQA.from_chain_type(
     llm=llm,
-    retriever=retriever,  
+    retriever=retriever,  # ارسال نمونه از retriever
     chain_type="stuff",
     input_key="question"
 )
-
 # ----------------- استیت برای چت -----------------
 if 'messages' not in st.session_state:
     st.session_state.messages = []