graphdoc.train package

class graphdoc.train.TrainerFactory

Bases: object

static single_trainer(trainer_class: str, prompt: SinglePrompt, optimizer_type: str, optimizer_kwargs: Dict[str, Any], mlflow_tracking_uri: str, mlflow_model_name: str, mlflow_experiment_name: str, trainset: List[Example], evalset: List[Example])

Returns an instance of the specified trainer class.

class graphdoc.train.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.

class graphdoc.train.DocQualityTrainer(prompt: DocQualityPrompt, 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: DocQualityPrompt, 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 DocQualityTrainer. This is the base class for implementing a trainer for a DocQualityPrompt.

Parameters:

• prompt (DocQualityPrompt) – 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.

evaluation_metrics(base_evaluation, optimized_evaluation)

Log evaluation metrics to mlflow. We will log the overall scores and the per category scores. Per category scores will be logged as a csv file.

Parameters:

• base_evaluation (Any) – The evaluation metrics of the base model.

• optimized_evaluation (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 provided, we will load the model from mlflow. Otherwise, we will use the provided DocQualityPrompt as the base model.

Parameters:

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

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

graphdoc.train.optimizer_class(optimizer_type: str, optimizer_kwargs: Dict[str, Any])

graphdoc.train.optimizer_compile(optimizer_type: str, optimizer_kwargs: Dict[str, Any])

Compiles the optimizer given the optimizer type and optimizer kwargs.

Optimizer kwargs are optimizer specific, and must include a student field that maps to a dspy.ChainOfThought, dspy.Predict, etc.

class graphdoc.train.SinglePrompt(prompt: Signature | SignatureMeta, prompt_type: Literal['predict', 'chain_of_thought'] | Callable, prompt_metric: Any)

Bases: ABC

__init__(prompt: Signature | SignatureMeta, prompt_type: Literal['predict', 'chain_of_thought'] | Callable, prompt_metric: Any) → None

Initialize a single prompt.

Parameters:

• prompt (dspy.Signature) – The prompt to use.

• prompt_type (Union[Literal["predict", "chain_of_thought"], Callable]) – The type of prompt to use. Can be “predict” or “chain_of_thought”. Optionally, pass another dspy.Module.

• prompt_metric (Any) – The metric to use. Marked as Any for flexibility (as metrics can be other prompts).

abstract evaluate_metric(example: Example, prediction: Prediction, trace=None) → Any

This is the metric used to evalaute the prompt.

Parameters:

• example (dspy.Example) – The example to evaluate the metric on.

• prediction (dspy.Prediction) – The prediction to evaluate the metric on.

• trace (Any) – The trace to evaluate the metric on. This is for DSPy.

abstract format_metric(examples: List[Example], overall_score: float, results: List, scores: List) → Dict[str, Any]

This takes the results from the evaluate_evalset and does any necessary formatting, taking into account the metric type.

Parameters:

• examples (List[dspy.Example]) – The examples to evaluate the metric on.

• overall_score (float) – The overall score of the metric.

• results (List) – The results from the evaluate_evalset.

• scores (List) – The scores from the evaluate_evalset.

abstract compare_metrics(base_metrics: Any, optimized_metrics: Any, comparison_value: str = 'overall_score') → bool

Compare the metrics of the base and optimized models. Return true if the optimized model is better than the base model.

Parameters:

• base_metrics (Any) – The metrics of the base model.

• optimized_metrics (Any) – The metrics of the optimized model.

• comparison_value (str) – The value to compare the metrics on. Determines which metric is used to compare the models.

Returns:

True if the optimized model is better than the base model, False otherwise.

Return type:

bool

evaluate_evalset(examples: List[Example], num_threads: int = 1, display_progress: bool = True, display_table: bool = True) → Dict[str, Any]

Take in a list of examples and evaluate the results.

Parameters:

• examples (List[dspy.Example]) – The examples to evaluate the results on.

• num_threads (int) – The number of threads to use for evaluation.

• display_progress (bool) – Whether to display the progress of the evaluation.

• display_table (bool) – Whether to display the table of the evaluation.

Returns:

A dictionary containing the overall score, results, and scores.

Return type:

Dict[str, Any]

class graphdoc.train.SinglePromptTrainer(prompt: SinglePrompt, 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: ABC

__init__(prompt: SinglePrompt, 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 SinglePromptTrainer. This is the base class for implementing a trainer for a single prompt.

Parameters:

• prompt (SinglePrompt) – 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.

abstract evaluation_metrics(base_evaluation, optimized_evaluation)

Log evaluation metrics to mlflow.

Parameters:

• base_evaluation (Any) – The evaluation metrics of the base model.

• optimized_evaluation (Any) – The evaluation metrics of the optimized model.

abstract 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.

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

Train the model.

Parameters:

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

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

Submodules

• graphdoc.train.doc_generator_trainer module
• graphdoc.train.doc_quality_trainer module
• graphdoc.train.optimizers module
• graphdoc.train.single_prompt_trainer module