Ayomide
25 lines
686 B
Python
25 lines
686 B
Python
import cohere
|
|
from .config import Config
|
|
|
|
|
|
class EmbeddingGenerator:
|
|
def __init__(self):
|
|
self.client = cohere.Client(Config.COHERE_API_KEY)
|
|
|
|
def generate_embeddings(self, text: str):
|
|
response = self.client.embed(
|
|
texts=[text],
|
|
model=Config.EMBED_MODEL,
|
|
input_type="document"
|
|
)
|
|
return response.embeddings[0]
|
|
|
|
def rerank_issues(self, issues: list, query: str, top_n: int = 5):
|
|
response = self.client.rerank(
|
|
query=query,
|
|
documents=issues,
|
|
top_n=top_n,
|
|
model=Config.RERANK_MODEL
|
|
)
|
|
return [result.document for result in response.results]
|