Spaces:
Build error
Build error
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]]) | |