biosummarize-ai / qa_engine.py
AryanRajSaxena's picture
second iteration
f4dcafb
raw
history blame contribute delete
628 Bytes
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
class QABot:
def __init__(self, chunks):
self.model = SentenceTransformer("pritamdeka/BioBERT-mnli-snli-scinli-scitail-mednli-stsb")
self.chunks = chunks
self.embeddings = self.model.encode(chunks)
self.index = faiss.IndexFlatL2(self.embeddings.shape[1])
self.index.add(np.array(self.embeddings))
def retrieve(self, query, k=3):
query_vec = self.model.encode([query])
D, I = self.index.search(np.array(query_vec), k)
return "\n\n".join([self.chunks[i] for i in I[0]])