Spaces:
Runtime error
Runtime error
File size: 1,089 Bytes
06bca0c 6aad21a 06bca0c 6aad21a 06bca0c |
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 |
from abc import ABC, abstractmethod
from dataclasses import dataclass
import pandas as pd
from openai.embeddings_utils import cosine_similarity
ALL_SOURCES = "All"
@dataclass
class Retriever(ABC):
@abstractmethod
def get_documents(self, source: str) -> pd.DataFrame:
"""Get all current documents from a given source."""
...
@abstractmethod
def get_source_display_name(self, source: str) -> str:
"""Get the display name of a source."""
...
def retrieve(self, query_embedding: list[float], top_k: int, source: str = None) -> pd.DataFrame:
documents = self.get_documents(source)
documents["similarity"] = documents.embedding.apply(lambda x: cosine_similarity(x, query_embedding))
# sort the matched_documents by score
matched_documents = documents.sort_values("similarity", ascending=False)
# limit search to top_k matched_documents.
top_k = len(matched_documents) if top_k == -1 else top_k
matched_documents = matched_documents.head(top_k)
return matched_documents
|