Spaces:
Runtime error
Runtime error
| from langchain.document_loaders.unstructured import UnstructuredFileLoader | |
| from langchain.text_splitter import CharacterTextSplitter | |
| from langchain.embeddings import OpenAIEmbeddings | |
| 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 transformers import AutoTokenizer, T5ForConditionalGeneration | |
| from retrieval.retrieval import Retrieval, BM25 | |
| import os, time | |
| class Agent: | |
| def __init__(self, args=None) -> None: | |
| self.args = args | |
| self.corpus = Retrieval() | |
| self.choices = args.choices | |
| self.context_value = "" | |
| self.use_context = False | |
| print("Model is loading...") | |
| self.model = T5ForConditionalGeneration.from_pretrained(args.model).to(args.device) | |
| self.tokenizer = AutoTokenizer.from_pretrained(args.tokenizer) | |
| print("Model loaded!") | |
| def load_context(self, doc_path): | |
| loader = UnstructuredFileLoader(doc_path.name) | |
| print('Loading file:', doc_path.name) | |
| context = loader.load()[0].page_content | |
| self.retrieval = Retrieval(docs=context) | |
| self.choices = self.retrieval.k | |
| self.use_context = True | |
| return f"Using file from {doc_path.name}" | |
| def asking(self, question): | |
| s_query = time.time() | |
| if self.use_context: | |
| print("Answering with your context") | |
| contexts = self.retrieval.get_context(question) | |
| else: | |
| print("Answering without your context") | |
| contexts = self.corpus.get_context(question) | |
| prompts = [] | |
| for context in contexts: | |
| prompt = f"Trả lời câu hỏi: {question} Trong nội dung: {context['context']}" | |
| prompts.append(prompt) | |
| s_token = time.time() | |
| tokens = self.tokenizer(prompts, max_length=self.args.seq_len, truncation=True, padding='max_length', return_tensors='pt') | |
| s_gen = time.time() | |
| 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 | |
| ) | |
| 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] | |
| 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 | |
| def get_context(self, context): | |
| self.context_value = context | |
| self.retrieval = Retrieval(docs=context) | |
| self.choices = self.retrieval.k | |
| self.use_context = True | |
| return context | |
| def load_context_file(self, file): | |
| print('Loading file:', file.name) | |
| text = '' | |
| for line in open(file.name, 'r', encoding='utf8'): | |
| text += line | |
| self.context_value = text | |
| return text | |
| def clear_context(self): | |
| self.context_value = "" | |
| self.use_context = False | |
| self.choices = self.args.choices | |
| return "" | |