| Current Path : /var/www/pythonian/streamlit/.venv/lib64/python3.10/site-packages/cohere/ |
| Current File : /var/www/pythonian/streamlit/.venv/lib64/python3.10/site-packages/cohere/classify.py |
from typing import Dict, List, NamedTuple
from cohere.response import CohereObject
LabelPrediction = NamedTuple("LabelPrediction", [("confidence", float)])
Example = NamedTuple("Example", [("text", str), ("label", str)])
class Classification(CohereObject):
def __init__(self, input: str, prediction: str, confidence: float, labels: Dict[str, LabelPrediction], *args,
**kwargs) -> None:
super().__init__(*args, **kwargs)
self.input = input
self.prediction = prediction
self.confidence = confidence
self.labels = labels
def __repr__(self) -> str:
prediction = self.prediction
confidence = self.confidence
labels = self.labels
return f"Classification<prediction: \"{prediction}\", confidence: {confidence}, labels: {labels}>"
class Classifications(CohereObject):
def __init__(self, classifications: List[Classification]) -> None:
self.classifications = classifications
def __repr__(self) -> str:
return self.classifications.__repr__()
def __str__(self) -> str:
return self.classifications.__str__()
def __iter__(self) -> iter:
return iter(self.classifications)
def __len__(self) -> int:
return len(self.classifications)
def __getitem__(self, index) -> Classification:
return self.classifications[index]