Spaces:
Build error
Build error
File size: 628 Bytes
7c114b1 f4dcafb 7c114b1 f4dcafb 7c114b1 f4dcafb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 |
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]])
|