Crizomb/ai_pdf
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

* added log tab

Browse source 
* added references text box
* added options to choose embedding models
This commit is contained in:
Crizomb 2024-04-20 12:54:24 +02:00
parent eedbb1b81a
commit 18f35b28c2
42 changed files with 911 additions and 441 deletions
Download patch file Download diff file Expand all files Collapse all files

26
backend/inference.py
View file

@ -15,22 +15,22 @@ class InferenceInstance:
self.vector_db_manager = vector_db_manager
self.history = []
self.nb_chunks_retrieved = nb_chunks_retrieved
flush_relevant_content()
def get_next_token(self, input_user: str, doc_name: str) -> Iterator[Dict[str, str]]:
is_pdf = doc_name.endswith(".pdf")
print(f"doc_name: {doc_name}")
new_assistant_message = {"role": "assistant", "content": ""}
search_results = self._get_search_results(input_user, doc_name)
self._update_history(input_user, search_results, is_pdf)
print(f"search results: {search_results}")
pages = self._update_history(input_user, search_results, is_pdf)
pages_info = f"pages used : p" + " p".join(pages)
print(f"history: {self.history}")
completion = self._get_completion()
for chunk in completion:
if chunk.choices[0].delta.content:
new_assistant_message["content"] += chunk.choices[0].delta.content
yield pages_info + "\n\n " + new_assistant_message["content"]
yield new_assistant_message["content"]
def _get_search_results(self, input_user: str, doc_name: str):
print(f"input_user: {input_user}")
@ -38,12 +38,21 @@ class InferenceInstance:
return vector_db.similarity_search(input_user, k=4)
def _update_history(self, input_user: str, search_results, is_pdf):
references_textbox_content = ""
some_context = ""
pages = []
for result in search_results:
if is_pdf:
pages.append(str(result.metadata['page']))
some_context += result.page_content + "\n\n"
pages_info = f'on page {result.metadata["page"]}' if is_pdf else 'in the document'
references_textbox_content += f"**Relevant content viewed {pages_info}**: \n\n" \
f" \n\n {result.page_content}\n\n" \
"-----------------------------------\n\n"
with open("../temp_file/relevant_content.mmd", "w") as f:
f.write(references_textbox_content)
self.history.append({"role": "system", "content": f"relevant content for user question {some_context}"})
self.history.append({"role": "user", "content": input_user})
return pages
@ -56,3 +65,12 @@ class InferenceInstance:
stream=True,
)
def read_relevant_content():
with open("../temp_file/relevant_content.mmd", "r") as f:
return f.read()
def flush_relevant_content():
with open("../temp_file/relevant_content.mmd", "w") as f:
f.write("")