Spaces:
Running
Running
| from typing import Union, List | |
| from datasets import Dataset | |
| from datasets.formatting.formatting import LazyRow | |
| from src.language_model.language_model_abstraction import LanguageModel | |
| from src.language_model.private_language_model_factory import ( | |
| private_language_model_factory, | |
| ) | |
| from src.task.task import Task, TaskType | |
| class RemoteLLMModel(LanguageModel): | |
| """ | |
| LLM Model based on private remote LLM provider (e.g. OpenAI) and pipeline mechanism for inference. | |
| """ | |
| def generate(self, rows: LazyRow) -> Union[str, List[str]]: | |
| generate = self.model.predict(rows["text"]) | |
| return generate | |
| def infer(self, rows: LazyRow) -> Union[str, List[str]]: | |
| infer = self.model.predict(rows["text"]) | |
| return infer | |
| def __init__( | |
| self, | |
| model_name: str, | |
| token: Union[str, None] = None, | |
| ): | |
| super().__init__(model_name) | |
| self._model_name = model_name | |
| self._token = token | |
| self.model = private_language_model_factory(model_name=self._model_name) | |
| def predict(self, evaluation_dataset: Dataset, task: Task) -> List: | |
| if task.task_type == TaskType.INFERENCE: | |
| labels = task.dataset.possible_ground_truths | |
| self.model.init_function_calling( | |
| labels, tool_choices=self.model.tool_choices | |
| ) | |
| inference_fn = self.infer | |
| else: | |
| inference_fn = self.generate | |
| process_dataset = evaluation_dataset.map( | |
| inference_fn, | |
| batched=False, | |
| desc=f"Running evaluation for task: {task.task_name}", | |
| remove_columns="text", | |
| ) | |
| self.model.print_none() | |
| return list(process_dataset["prediction"]) | |