Autonomous-AI / knowledge.py
Leonydis137's picture
Upload 2 files
d5499fd verified
raw
history blame
851 Bytes
from sentence_transformers import SentenceTransformer
import faiss
import os
class KnowledgeIngestor:
def __init__(self):
self.model = SentenceTransformer("all-MiniLM-L6-v2")
self.index = faiss.IndexFlatL2(384)
self.texts = []
def ingest_directory(self, dir_path):
for fname in os.listdir(dir_path):
if fname.endswith(".txt"):
with open(os.path.join(dir_path, fname), "r", encoding="utf-8") as f:
content = f.read()
self.add_text(content)
def add_text(self, text):
vec = self.model.encode([text])
self.index.add(vec)
self.texts.append(text)
def search(self, query, top_k=3):
vec = self.model.encode([query])
D, I = self.index.search(vec, top_k)
return [self.texts[i] for i in I[0]]