seai2526-uniba-TheClouds/Code-Comment-Classification-Api

1"""Module for training different types of models for code comment classification."""
2
3import argparse
4import logging
5import os
6
7import dagshub
8from datasets import Dataset
9import mlflow
10import yaml
11
12from .utils import load_dataset_splits, parse_labels_column
13
14logging.basicConfig(
15    level=logging.INFO,
16    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
17)
18logger = logging.getLogger(__name__)
19
20
21dagshub.init(repo_owner="se4ai2526-uniba", repo_name="TheClouds", mlflow=True)
22
23
24def train_model(lang, model_type, data_path, model_output_path, params):
25    """Trains and saves a model for a specific language and model type."""
26    print(f"--- Starting training for language: {lang} with model: {model_type} ---")
27
28    ds = load_dataset_splits(data_path)
29
30    train_df = ds[f"{lang}_train"]
31    eval_df = ds[f"{lang}_test"]
32
33    train_df = parse_labels_column(train_df)
34    eval_df = parse_labels_column(eval_df)
35
36    # converto i DataFrame in HuggingFace Dataset
37    train_dataset = Dataset.from_pandas(train_df, preserve_index=False)
38    eval_dataset = Dataset.from_pandas(eval_df, preserve_index=False)
39
40    if model_type == "setfit":
41        from setfit import SetFitModel, Trainer, TrainingArguments
42
43        mlflow.set_experiment("SetFit Training")
44        with mlflow.start_run(run_name=f"train-{lang}-{model_type}"):
45            mlflow.log_param("language", lang)
46            mlflow.log_param("model_type", model_type)
47            model = SetFitModel.from_pretrained(
48                "sentence-transformers/paraphrase-MiniLM-L6-v2",
49                multi_target_strategy="multi-output",
50            )
51            args = TrainingArguments(**params)
52            trainer = Trainer(
53                model=model,
54                args=args,
55                train_dataset=train_dataset,
56                eval_dataset=eval_dataset,
57                column_mapping={"combo": "text", "labels": "label"},
58            )
59
60            mlflow.log_param("num_epochs", args.num_epochs)
61            mlflow.log_param("num_iterations", args.num_iterations)
62
63            trainer.train()
64
65            eval_metrics = trainer.evaluate()
66            for metric_name, metric_value in eval_metrics.items():
67                mlflow.log_metric(metric_name, metric_value)
68
69            trainer.model.save_pretrained(model_output_path)
70
71            mlflow.transformers.log_model(
72                transformers_model=model_output_path,
73                artifact_path=f"{lang}_setfit_model",
74                task="text-classification",
75            )
76            mlflow.end_run()
77
78    elif model_type == "random_forest":
79        import joblib
80        import numpy as np
81        from sklearn.ensemble import RandomForestClassifier
82        from sklearn.feature_extraction.text import TfidfVectorizer
83        from sklearn.multioutput import MultiOutputClassifier
84        from sklearn.pipeline import Pipeline
85
86        mlflow.set_experiment("Random Forest Training")
87        with mlflow.start_run(run_name=f"train-{lang}-{model_type}"):
88            mlflow.log_param("language", lang)
89            mlflow.log_param("model_type", model_type)
90            mlflow.log_params(params)
91
92            tfidf_params = {
93                "ngram_range": tuple(params.pop("ngram_range", (1, 1))),
94                "max_features": params.pop("max_features", None),
95                "min_df": params.pop("min_df", 1),
96                "max_df": params.pop("max_df", 1.0),
97            }
98
99            rf_params = params
100            pipeline = Pipeline(
101                [
102                    ("tfidf", TfidfVectorizer(**tfidf_params)),
103                    (
104                        "clf",
105                        MultiOutputClassifier(
106                            RandomForestClassifier(
107                                random_state=42, class_weight="balanced", **rf_params
108                            )
109                        ),
110                    ),
111                ]
112            )
113
114            X_train = train_dataset["combo"]
115            y_train = np.array(train_dataset["labels"])
116
117            pipeline.fit(X_train, y_train)
118
119            X_test = eval_dataset["combo"]
120            y_test = np.array(eval_dataset["labels"])
121
122            score = pipeline.score(X_test, y_test)
123            mlflow.log_metric("accuracy", score)
124
125            os.makedirs(os.path.dirname(model_output_path), exist_ok=True)
126            joblib.dump(pipeline, f"{model_output_path}.joblib")
127
128            mlflow.sklearn.log_model(
129                sk_model=pipeline, artifact_path=f"{lang}_random_forest_model"
130            )
131            mlflow.end_run()
132
133    elif model_type == "transformer":
134        from .transformer import (
135            TransformerConfig,
136            TransformerTrainer,
137        )
138
139        mlflow.set_experiment("Transformer Training")
140        with mlflow.start_run(run_name=f"train-{lang}-{model_type}"):
141            mlflow.log_param("language", lang)
142            mlflow.log_param("model_type", model_type)
143            mlflow.log_params(params)
144
145            cfg = TransformerConfig(
146                lang=lang,
147                raw_data_dir="data/raw",
148                processed_data_dir="data/processed/transformer",
149                model_output_path=model_output_path,
150                pretrained_model_name=params.get(
151                    "pretrained_model_name", "microsoft/codebert-base"
152                ),
153                max_length=params.get("max_length", 128),
154                batch_size=params.get("batch_size", 16),
155                lr=params.get("lr", 2e-5),
156                num_epochs=params.get("num_epochs", 5),
157                warmup_ratio=params.get("warmup_ratio", 0.1),
158                pos_weight_cap=params.get("pos_weight_cap", 30.0),
159                threshold=params.get("threshold", 0.5),
160                preprocessing=params.get("preprocessing", False),
161                preprocessing_factor=params.get("preprocessing_factor", 1.0),
162            )
163
164            logger.info(
165                "Starting transformer training for language '%s' with config: %s",
166                lang,
167                cfg,
168            )
169
170            trainer = TransformerTrainer(cfg)
171            metrics = trainer.run()
172
173            logger.info("Final transformer metrics for %s: %s", lang, metrics)
174
175            for name, value in metrics.items():
176                mlflow.log_metric(f"final_{name}", value)
177
178            mlflow.end_run()
179
180    else:
181        raise ValueError(f"Unsupported model_type: {model_type}")
182
183    print(f"Model for {lang}-{model_type} saved to {model_output_path}")
184
185
186if __name__ == "__main__":
187    parser = argparse.ArgumentParser()
188    parser.add_argument("--lang", type=str, required=True)
189    parser.add_argument("--model_type", type=str, required=True)
190    args = parser.parse_args()
191
192    with open("params.yaml", "r") as f:
193        all_params = yaml.safe_load(f)
194
195    model_params = all_params[args.model_type].copy()
196
197    train_model(
198        lang=args.lang,
199        model_type=args.model_type,
200        data_path="data/raw",
201        model_output_path=f"models/{args.lang}/{args.model_type}",
202        params=model_params,
203    )
204