typing fix

Author: finitearth

promptolution/helpers.py

@@ -161,9 +161,9 @@ def get_llm(model_id: Optional[str] = None, config: Optional["ExperimentConfig"]
 def get_task(
     df: pd.DataFrame,
     config: "ExperimentConfig",
-    task_type: TaskType = None,
-    judge_llm: "BaseLLM" = None,
-    reward_function: Callable = None,
+    task_type: Optional["TaskType"] = None,
+    judge_llm: Optional["BaseLLM"] = None,
+    reward_function: Optional[Callable] = None,
 ) -> "BaseTask":
     """Get the task based on the provided DataFrame and configuration.
 
@@ -198,7 +198,7 @@ def get_optimizer(
     predictor: "BasePredictor",
     meta_llm: "BaseLLM",
     task: "BaseTask",
-    optimizer: Optional[OptimizerType] = None,
+    optimizer: Optional["OptimizerType"] = None,
     task_description: Optional[str] = None,
     config: Optional["ExperimentConfig"] = None,
 ) -> "BaseOptimizer":
@@ -292,7 +292,7 @@ def get_exemplar_selector(
         raise ValueError(f"Unknown exemplar selector: {name}")
 
 
-def get_predictor(downstream_llm=None, type: PredictorType = "marker", *args, **kwargs) -> "BasePredictor":
+def get_predictor(downstream_llm=None, type: "PredictorType" = "marker", *args, **kwargs) -> "BasePredictor":
     """Factory function to create and return a predictor instance.
 
     This function supports three types of predictors:

promptolution/tasks/base_task.py

@@ -27,7 +27,7 @@ def __init__(
         y_column: Optional[str] = None,
         task_description: Optional[str] = None,
         n_subsamples: int = 30,
-        eval_strategy: EvalStrategy = "full",
+        eval_strategy: "EvalStrategy" = "full",
         seed: int = 42,
         config: Optional["ExperimentConfig"] = None,
     ) -> None:
@@ -70,7 +70,7 @@ def __init__(
         self.eval_cache: Dict[Tuple[str, str, str], float] = {}  # (prompt, x, y): scores per datapoint
         self.seq_cache: Dict[Tuple[str, str, str], str] = {}  # (prompt, x, y): generating sequence per datapoint
 
-    def subsample(self, eval_strategy: EvalStrategy = None) -> Tuple[List[str], List[str]]:
+    def subsample(self, eval_strategy: "EvalStrategy" = None) -> Tuple[List[str], List[str]]:
         """Subsample the dataset based on the specified parameters.
 
         Args:
@@ -170,7 +170,7 @@ def evaluate(
         system_prompts: Optional[Union[str, List[str]]] = None,
         return_agg_scores: Literal[True] = True,
         return_seq: Literal[False] = False,
-        eval_strategy: Optional[EvalStrategy] = None,
+        eval_strategy: Optional["EvalStrategy"] = None,
     ) -> List[float]:
         ...
 
@@ -182,7 +182,7 @@ def evaluate(
         system_prompts: Optional[Union[str, List[str]]] = None,
         return_agg_scores: Literal[False] = False,
         return_seq: Literal[False] = False,
-        eval_strategy: Optional[EvalStrategy] = None,
+        eval_strategy: Optional["EvalStrategy"] = None,
     ) -> List[List[float]]:
         ...
 
@@ -194,7 +194,7 @@ def evaluate(
         system_prompts: Optional[Union[str, List[str]]] = None,
         return_agg_scores: Literal[False] = False,
         return_seq: Literal[True] = True,
-        eval_strategy: Optional[EvalStrategy] = None,
+        eval_strategy: Optional["EvalStrategy"] = None,
     ) -> Tuple[List[List[float]], List[List[str]]]:
         ...
 
@@ -206,7 +206,7 @@ def evaluate(
         system_prompts: Optional[Union[str, List[str]]] = None,
         return_agg_scores: Literal[True] = True,
         return_seq: Literal[False] = False,
-        eval_strategy: Optional[EvalStrategy] = None,
+        eval_strategy: Optional["EvalStrategy"] = None,
     ) -> List[float]:
         ...
 
@@ -218,7 +218,7 @@ def evaluate(
         system_prompts: Optional[Union[str, List[str]]] = None,
         return_agg_scores: Literal[False] = False,
         return_seq: Literal[False] = False,
-        eval_strategy: Optional[EvalStrategy] = None,
+        eval_strategy: Optional["EvalStrategy"] = None,
     ) -> List[List[float]]:
         ...
 
@@ -230,7 +230,7 @@ def evaluate(
         system_prompts: Optional[Union[str, List[str]]] = None,
         return_agg_scores: Literal[False] = False,
         return_seq: Literal[True] = True,
-        eval_strategy: Optional[EvalStrategy] = None,
+        eval_strategy: Optional["EvalStrategy"] = None,
     ) -> Tuple[List[List[float]], List[List[str]]]:
         ...
 
@@ -241,7 +241,7 @@ def evaluate(
         system_prompts: Optional[Union[str, List[str]]] = None,
         return_agg_scores: bool = True,
         return_seq: bool = False,
-        eval_strategy: Optional[EvalStrategy] = None,
+        eval_strategy: Optional["EvalStrategy"] = None,
     ) -> Union[List[float], List[List[float]], Tuple[List[List[float]], List[List[str]]]]:
         """Evaluate a set of prompts using a given predictor.
 

promptolution/tasks/judge_tasks.py

@@ -71,7 +71,7 @@ def __init__(
         y_column: Optional[str] = None,
         task_description: Optional[str] = None,
         n_subsamples: int = 30,
-        eval_strategy: EvalStrategy = "full",
+        eval_strategy: "EvalStrategy" = "full",
         seed: int = 42,
         judge_prompt: Optional[str] = None,
         min_score: float = -5.0,

promptolution/tasks/reward_tasks.py

@@ -1,10 +1,9 @@
 """Module for Reward tasks."""
 
 
-import numpy as np
 import pandas as pd
 
-from typing import TYPE_CHECKING, Callable, List, Literal, Optional, Union
+from typing import TYPE_CHECKING, Callable, List, Optional
 
 from promptolution.tasks.base_task import BaseTask
 
@@ -27,7 +26,7 @@ def __init__(
         x_column: str = "x",
         task_description: Optional[str] = None,
         n_subsamples: int = 30,
-        eval_strategy: EvalStrategy = "full",
+        eval_strategy: "EvalStrategy" = "full",
         seed: int = 42,
         config: Optional["ExperimentConfig"] = None,
     ) -> None: