rasmus1610's picture
reranking and polishing
e22d4b7
raw
history blame
3.88 kB
import numpy as np
import pandas as pd
from rank_bm25 import BM25Okapi
from llm import OpenAILLM
def cosine_similarity(vector1, vector2):
return np.dot(vector1, vector2) / (np.linalg.norm(vector1) * np.linalg.norm(vector2))
class QuestionAnswerer:
def __init__(self, docs, embedding_model, llm=OpenAILLM('gpt-3.5-turbo'), cross_encoder=None):
self.docs = docs
self.bm25 = BM25Okapi([c.split(" ") for c in self.docs.chunks.values[1:]])
self.embedding_model = embedding_model
self.llm = llm
self.cross_encoder = cross_encoder
def sim_search(self, query, n=10, use_hyde=False, use_dot_product=False):
if use_hyde:
generated_doc = self._get_generated_doc(query)
print("generated document (hyde): \n", generated_doc)
embedding = self.embedding_model.encode(generated_doc)
else:
embedding = self.embedding_model.encode(query)
if use_dot_product:
similarities = self.docs['embeddings'].apply(lambda x: np.dot(x, embedding))
else:
similarities = self.docs['embeddings'].apply(lambda x: cosine_similarity(x, embedding))
self.docs['similarities'] = similarities
return self.docs.sort_values('similarities', ascending=False).head(n)
def sim_search_rerank(self, query, n=10, sim_search_n=100, **kwargs):
search_results = self.sim_search(query, n=sim_search_n, use_hyde=False, **kwargs)
reranked_results = self.rerank(search_results, query)
return reranked_results[:n]
def sim_search_bm25(self, query, n=10):
tokenized_query = query.split(" ")
doc_scores = self.bm25.get_scores(tokenized_query)
self.docs['bm25'] = np.insert(doc_scores, 0, 0) #hack because I have to remove the first item, because I cannot split it
result = self.docs.sort_values('bm25', ascending=False)[:n]
return result
def _create_prompt(self, context, question):
return f"""
Context information is below.
---------------------
{context}
---------------------
Given the context information and not prior knowledge, answer the query.
Query: {question}
Answer: \
"""
def _get_generated_doc(self, question):
prompt = f"""Write a guideline section in German answering the question below
---------------------
Question: {question}
---------------------
Answer: \
"""
system_prompt = "You are an experienced radiologist answering medical questions. You answer only in German."
return self.llm.get_response(system_prompt, prompt)
def rerank(self, docs, query):
inp = [[query, doc.chunks] for doc in docs.itertuples()]
cross_scores = self.cross_encoder.predict(inp) if self.cross_encoder else []
docs['cross_score'] = cross_scores
return docs.sort_values('cross_score', ascending=False)
def answer_question(self, question, n=3, use_hyde=False, use_reranker=False, use_bm25=False):
if use_reranker and use_hyde:
print('reranking together with hyde is not supported yet')
if use_reranker:
search_results = self.sim_search_rerank(question, n=n)
if use_bm25:
search_results = self.sim_search_bm25(question, n=n)
else:
search_results = self.sim_search(question, n=n, use_hyde=use_hyde)
context = "\n\n".join(search_results['chunks'])
prompt = self._create_prompt(context, question)
system_prompt = "You are a helpful assistant answering questions in German. You answer only in German. If you do not know an answer you say it. You do not fabricate answers."
return self.llm.get_response(system_prompt, prompt)