Skip to content

Instantly share code, notes, and snippets.

@fabian57fabian
Last active April 3, 2024 15:33
Show Gist options
  • Save fabian57fabian/5b4d67728277d3800e98995bef0bb764 to your computer and use it in GitHub Desktop.
Save fabian57fabian/5b4d67728277d3800e98995bef0bb764 to your computer and use it in GitHub Desktop.
deepeval class for cohere LLM api (command-r model)
from typing import Optional
from cohere import Client
from deepeval.models.base_model import DeepEvalBaseLLM
class CohereModel(DeepEvalBaseLLM):
def __init__(
self,
cohere_api_key: str,
model_name: Optional[str] = "command-r",
max_tokens: Optional[int] = 1024,
temperature: Optional[float] = 0.7,
):
self.cohere_api_key = cohere_api_key
self.model_name = model_name
self.max_tokens = max_tokens
self.temperature = temperature
super().__init__()
def load_model(self):
return Client(self.cohere_api_key)
def generate(self, prompt: str) -> str:
chat_model = self.load_model()
return chat_model.chat(message=prompt,
max_tokens=self.max_tokens,
temperature=self.temperature).text
async def a_generate(self, prompt: str) -> str:
chat_model = self.load_model()
res = await chat_model.chat(message=prompt,
max_tokens=self.max_tokens,
temperature=self.temperature)
return res.text
def get_model_name(self):
return self.model_name
@fabian57fabian
Copy link
Author

fabian57fabian commented Apr 3, 2024

Usage (Example with summarization metric):

from deepeval.metrics import SummarizationMetric
from cohere.model import CohereModel

api_key_cohere = "..."
model = CohereModel(api_key_cohere)
metric_using_cohere_llm = SummarizationMetric(model=model)

# use metric_using_cohere_llm in your test case

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment