graphdoc.train.doc_generator_trainer module

class graphdoc.train.doc_generator_trainer.DocGeneratorTrainer(prompt: DocGeneratorPrompt, optimizer_type: str, optimizer_kwargs: Dict[str, Any], mlflow_model_name: str, mlflow_experiment_name: str, mlflow_tracking_uri: str, trainset: List[Example], evalset: List[Example])

Bases: SinglePromptTrainer

__init__(prompt: DocGeneratorPrompt, optimizer_type: str, optimizer_kwargs: Dict[str, Any], mlflow_model_name: str, mlflow_experiment_name: str, mlflow_tracking_uri: str, trainset: List[Example], evalset: List[Example])

Initialize the DocGeneratorTrainer.

Parameters:

• prompt (DocGeneratorPrompt) – The prompt to train.

• optimizer_type (str) – The type of optimizer to use.

• optimizer_kwargs (Dict[str, Any]) – The keyword arguments for the optimizer.

• mlflow_model_name (str) – The name of the model in mlflow.

• mlflow_experiment_name (str) – The name of the experiment in mlflow.

• mlflow_tracking_uri (str) – The uri of the mlflow tracking server.

• trainset (List[dspy.Example]) – The training set.

• evalset (List[dspy.Example]) – The evaluation set.

_calculate_average_score(evaluation: dict) → float

Given a dictionary of evaluation results, calculate the average score.

Parameters:

evaluation (Dict[str, Any]) – The evaluation results.

Returns:

The average score.

Return type:

float

evaluation_metrics(base_evaluation: Dict[str, Any], optimized_evaluation: Dict[str, Any]) → None

Log evaluation metrics to mlflow.

Parameters:

• base_evaluation (Dict[str, Any]) – The evaluation metrics of the base model.

• optimized_evaluation (Dict[str, Any]) – The evaluation metrics of the optimized model.

evaluate_training(base_model, optimized_model) → Tuple[Dict[str, Any], Dict[str, Any]]

Evaluate the training of the model. Comparing the base and optimized models.

Parameters:

• base_model (Any) – The base model.

• optimized_model (Any) – The optimized model.

train(load_model_args: Dict[str, Any] | None = None, save_model: bool = True)

Train the model. If load_model_args is provided, load the model from MLFlow.

Parameters:

• load_model_args (Optional[Dict[str, Any]]) – The arguments to load the model from mlflow.

• save_model (bool) – Whether to save the model to mlflow.