.

app.py

@@ -1,11 +1,60 @@
 import os
-import pandas as pd
+import torch
 import gradio as gr
-from persian_rag import PersianRAG
+import pandas as pd
+from sentence_transformers import SentenceTransformer, util
+from transformers import GenerationConfig, AutoModelForCausalLM, AutoTokenizer
 
 os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
 
 
+class PersianRAG:
+    def __init__(self, knowledge, embedding_model='LABSE', llm_model="MehdiHosseiniMoghadam/AVA-Mistral-7B-V2",
+                 device='cuda', retrieved_docs=3):
+        self.device = device
+        self.retrieved_docs = retrieved_docs
+        self.answer_df = (knowledge['Answer'])
+
+        self.embedder = SentenceTransformer(embedding_model)
+        self.question_embeddings = self.embedder.encode((knowledge['Question']), show_progress_bar=True,
+                                                        convert_to_tensor=True)
+
+        self.model = AutoModelForCausalLM.from_pretrained(llm_model, torch_dtype=torch.float16, device_map="auto")
+        self.tokenizer = AutoTokenizer.from_pretrained(llm_model)
+        self.generation_config = GenerationConfig(
+            do_sample=True,
+            top_k=1,
+            temperature=0.99,
+            max_new_tokens=900,
+            pad_token_id=self.tokenizer.eos_token_id
+        )
+
+    def rag(self, query):
+        ans = {}
+        question_embedding = self.embedder.encode(query, convert_to_tensor=True)
+        hits = util.semantic_search(question_embedding, self.question_embeddings)
+        hits = hits[0]
+
+        for hit in hits[0:self.retrieved_docs]:
+            ans[hit['corpus_id']] = self.answer_df[hit['corpus_id']]
+        ans = pd.DataFrame(list(ans.items()), columns=['id', 'res'])
+
+        prompt = f'''
+        با توجه به شرایط زیر به این سوال پاسخ دهید:
+        {query},
+
+        متن نوشته:
+        {ans['res'][0]} - {ans['res'][1]} - {ans['res'][2]}
+        '''
+
+        prompt = f"### Human:{prompt}\n### Assistant:"
+
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+        outputs = self.model.generate(**inputs, generation_config=self.generation_config)
+
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+
 # Function to load CSV and initialize PersianRAG
 def init_rag(knowledge_file, embedding_model, llm_model, device, retrieved_docs):
     knowledge = pd.read_csv(knowledge_file)

requirements.txt

@@ -0,0 +1,12 @@
+sentence-transformers
+metapub
+openpyxl
+accelerate
+transformers
+accelerate
+bitsandbytes
+trl
+py7zr
+optimum
+git+https://github.com/huggingface/peft
+# pip install auto-gptq --extra-index-url https://huggingface.github.io/autogptq-index/whl/cu118