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]])