Update app.py
Browse files
app.py
CHANGED
@@ -67,7 +67,7 @@ class Retriever:
|
|
67 |
def load_chunks(self):
|
68 |
self.text = self.extract_text_from_pdf(self.file_path)
|
69 |
text_splitter = RecursiveCharacterTextSplitter(
|
70 |
-
chunk_size=
|
71 |
chunk_overlap=20,
|
72 |
length_function=self.token_len,
|
73 |
separators=["Section", "\n\n", "\n", ".", " ", ""]
|
@@ -76,7 +76,7 @@ class Retriever:
|
|
76 |
self.chunks = text_splitter.split_text(self.text)
|
77 |
|
78 |
def load_context_embeddings(self):
|
79 |
-
encoded_input = self.context_tokenizer(self.chunks, return_tensors='pt', padding=True, truncation=True, max_length=
|
80 |
|
81 |
with torch.no_grad():
|
82 |
model_output = self.context_model(**encoded_input)
|
@@ -86,7 +86,7 @@ class Retriever:
|
|
86 |
self.index.add(self.token_embeddings)
|
87 |
|
88 |
def retrieve_top_k(self, query_prompt, k=10):
|
89 |
-
encoded_query = self.question_tokenizer(query_prompt, return_tensors="pt", max_length=
|
90 |
|
91 |
with torch.no_grad():
|
92 |
model_output = self.question_model(**encoded_query)
|
@@ -127,8 +127,8 @@ class RAG:
|
|
127 |
|
128 |
input_text = "answer: " + " ".join(context) + " " + question
|
129 |
|
130 |
-
inputs = self.generator_tokenizer.encode(input_text, return_tensors='pt', max_length=
|
131 |
-
outputs = self.generator_model.generate(inputs, max_length=
|
132 |
|
133 |
answer = self.generator_tokenizer.decode(outputs[0], skip_special_tokens=True)
|
134 |
return answer
|
|
|
67 |
def load_chunks(self):
|
68 |
self.text = self.extract_text_from_pdf(self.file_path)
|
69 |
text_splitter = RecursiveCharacterTextSplitter(
|
70 |
+
chunk_size=300,
|
71 |
chunk_overlap=20,
|
72 |
length_function=self.token_len,
|
73 |
separators=["Section", "\n\n", "\n", ".", " ", ""]
|
|
|
76 |
self.chunks = text_splitter.split_text(self.text)
|
77 |
|
78 |
def load_context_embeddings(self):
|
79 |
+
encoded_input = self.context_tokenizer(self.chunks, return_tensors='pt', padding=True, truncation=True, max_length=300).to(device)
|
80 |
|
81 |
with torch.no_grad():
|
82 |
model_output = self.context_model(**encoded_input)
|
|
|
86 |
self.index.add(self.token_embeddings)
|
87 |
|
88 |
def retrieve_top_k(self, query_prompt, k=10):
|
89 |
+
encoded_query = self.question_tokenizer(query_prompt, return_tensors="pt", max_length=300, truncation=True, padding=True).to(device)
|
90 |
|
91 |
with torch.no_grad():
|
92 |
model_output = self.question_model(**encoded_query)
|
|
|
127 |
|
128 |
input_text = "answer: " + " ".join(context) + " " + question
|
129 |
|
130 |
+
inputs = self.generator_tokenizer.encode(input_text, return_tensors='pt', max_length=300, truncation=True).to(device)
|
131 |
+
outputs = self.generator_model.generate(inputs, max_length=300, min_length=2, length_penalty=2.0, num_beams=4, early_stopping=True)
|
132 |
|
133 |
answer = self.generator_tokenizer.decode(outputs[0], skip_special_tokens=True)
|
134 |
return answer
|