Spaces:
No application file
No application file
File size: 1,059 Bytes
5a1b165 bcd324a 5a1b165 64ef551 5a1b165 64ef551 5a1b165 64ef551 5a1b165 fa46f05 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
from abc import ABC, abstractmethod
import jsonlines
from TwitterChatBot.utils import dot_similarity
class Index(ABC):
@abstractmethod
def load(self, path):
pass
class JsonLinesIndex(Index):
def load(self, path):
with jsonlines.open(path) as passages:
indexes = list(passages)
return indexes
class IndexSearchEngine:
def __init__(self, indexes, gpt_manager):
self.indexes = indexes
self.gpt_manager = gpt_manager
def search(self, question, count=4):
question_embedding = self.gpt_manager.get_embedding(prompt=question)
simmilarities = []
for index in self.indexes:
embedding = index["embedding"]
score = dot_similarity(question_embedding, embedding)
simmilarities.append({"index": index, "score": score})
sorted_similarities = sorted(
simmilarities, key=lambda x: x["score"], reverse=True
)
return [result["index"]["content"] for result in sorted_similarities[:count]]
|