Update agent_t5.py

agent_t5.py

@@ -5,21 +5,26 @@ from langchain.vectorstores import Chroma
 from langchain.chains import RetrievalQA
 from langchain.chat_models import ChatOpenAI
 from langchain.schema import AIMessage, HumanMessage, SystemMessage, Document
+from langchain.document_loaders import PyPDFLoader
 
 from transformers import AutoTokenizer, T5ForConditionalGeneration
 from retrieval.retrieval import Retrieval, BM25
-import os, time
+import os, time, torch
+from torch.nn import Softmax
 
 
 
 class Agent:
     def __init__(self, args=None) -> None:
         self.args = args
-        self.corpus = Retrieval()
         self.choices = args.choices
+        self.corpus = Retrieval(k=args.choices)
         
         self.context_value = ""
         self.use_context = False
+        self.softmax = Softmax(dim=1)
+        self.temp = []
+        self.replace_list = torch.load('retrieval/replace.pt')
         
         print("Model is loading...")
         self.model = T5ForConditionalGeneration.from_pretrained(args.model).to(args.device)
@@ -28,9 +33,12 @@ class Agent:
 
 
     def load_context(self, doc_path):
-        loader = UnstructuredFileLoader(doc_path.name)
         print('Loading file:', doc_path.name)
-        context = loader.load()[0].page_content
+        if doc_path.name[-4:] == '.pdf':
+            context = self.read_pdf(doc_path.name)
+        else:
+            # loader = UnstructuredFileLoader(doc_path.name)
+            context = open(doc_path.name, encoding='utf-8').read()
         
         self.retrieval = Retrieval(docs=context)
         self.choices = self.retrieval.k
@@ -42,10 +50,10 @@ class Agent:
     def asking(self, question):
         s_query = time.time()
         if self.use_context:
-            print("Answering with your context")
+            print("Answering with your context:", question)
             contexts = self.retrieval.get_context(question)
         else:
-            print("Answering without your context")
+            print("Answering without your context:", question)
             contexts = self.corpus.get_context(question)
 
         prompts = []
@@ -60,28 +68,48 @@ class Agent:
         outputs = self.model.generate(
             input_ids=tokens.input_ids.to(self.args.device),
             attention_mask=tokens.attention_mask.to(self.args.device),
-            max_new_tokens=self.args.out_len
+            max_new_tokens=self.args.out_len,
+            output_scores=True,
+            return_dict_in_generate=True
         )
 
-        s_de = time.time()
-        answers = []
-        for output in outputs:
-            sequence = self.tokenizer.decode(output, skip_special_tokens=True)
-            answers.append(sequence)
 
-        self.temp = [contexts, answers]
+        s_de = time.time()
+        results = []
+
+        scores = self.softmax(outputs.scores[0])
+        scores = scores.max(dim=1).values*100
+        # print(scores)
+        for i in range(self.choices):
+            result = contexts[i]
+            score = round(scores[i].item())
+            result['score'] = score
+            
+            answer = self.tokenizer.decode(outputs.sequences[i], skip_special_tokens=True)
+            result['answer'] = answer
+            results.append(result)
+
+        def get_score(record):
+            return record['score']**2 * record['score_bm']
+
+        results.sort(key=get_score, reverse=True)
+
+        self.temp = results
         t_mess = "t_query: {:.2f}\t t_token: {:.2f}\t t_gen: {:.2f}\t t_decode: {:.2f}\t".format(
             s_token-s_query, s_gen-s_token, s_de-s_gen, time.time()-s_de
         )
-        print(t_mess)
-        return answers
+        print(t_mess, len(self.temp))
+        if results[0]['score'] > 50:
+            return results[0]['answer']
+        else:
+            return f"Tôi không chắc nhưng câu trả lời có thể là: {results[0]['answer']}\nBạn có thể tham khảo các câu trả lời bên cạnh!"
 
 
     
     def get_context(self, context):
         self.context_value = context
 
-        self.retrieval = Retrieval(docs=context)
+        self.retrieval = Retrieval(k=self.choices, docs=context)
         self.choices = self.retrieval.k
         self.use_context = True
         return context
@@ -100,3 +128,19 @@ class Agent:
         self.use_context = False
         self.choices = self.args.choices
         return ""
+
+    def replace(self, text):
+        for key, value in self.replace_list:
+            text = text.replace(key, value)
+        return text
+
+    def read_pdf(self, file_path):
+        loader = PyPDFLoader(file_path)
+        pages = loader.load_and_split()
+        text = ''
+        for page in pages:
+            page_content = page.page_content
+            text += self.replace(page_content)
+
+        return text
+