import os import random import shutil from datasets import Dataset import joblib from loguru import logger import mlflow import numpy as np from numpy import ndarray from peft import LoraConfig, TaskType, get_peft_model from sentence_transformers import SentenceTransformer from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score from sklearn.multioutput import MultiOutputClassifier import torch import torch.nn as nn from transformers import get_linear_schedule_with_warmup from xgboost import XGBClassifier from turing.modeling.baseModel import BaseModel from turing.modeling.models.MiniLMClassifierWrapper import MiniLMClassifierWrapper def drop_tokens(text, drop_prob=0.1): """ Randomly drops tokens from the input text based on the specified drop probability. """ tokens = text.split() if len(tokens) <= 3: return text return " ".join( t for t in tokens if random.random() > drop_prob ) def drop_tokens_batch(texts, drop_prob=0.1, apply_prob=0.3): """ Apply token dropping augmentation to a batch of texts. """ augmented = [] for text in texts: if random.random() < apply_prob: augmented.append(drop_tokens(text, drop_prob)) elif random.random() < 0.15: x = " ".join(text.split()) augmented.append(x) else: augmented.append(text) return augmented def finetune_miniLM(X_train, y_train, device,model_save_path="sentence-transformers/minilm.pt"): """ Train MiniLM model with temporary classification head using java dataset only. Args: X_train: Input training data. y_train: True labels for training data. """ encoder = SentenceTransformer('sentence-transformers/paraphrase-MiniLM-L6-v2').to(device) peft_config = LoraConfig( task_type=TaskType.FEATURE_EXTRACTION, lora_alpha=16, bias="none", lora_dropout=0.1, ) encoder[0].auto_model = get_peft_model( encoder[0].auto_model, peft_config ) encoder[0].auto_model.print_trainable_parameters() y_train = np.array(y_train,dtype=np.float32) dataset = Dataset.from_dict({"text": X_train, "labels": y_train}) split_set = dataset.train_test_split(test_size= 0.2, seed=42) train_set = split_set['train'] eval_set = split_set['test'] epoch = 10 batch_size = 32 total_steps = len(train_set) // batch_size * epoch warm_up_steps = int(0.1 * total_steps) y_train = np.array(y_train,dtype=np.float32) classifier = nn.Sequential( nn.Linear(encoder.get_sentence_embedding_dimension(), 128), nn.ReLU(), nn.Dropout(0.1), nn.Linear(128, len(y_train[0])) ).to(device) logger.info(f"Training set size: {len(train_set)}, Evaluation set size: {len(eval_set)}") criterion = nn.BCEWithLogitsLoss() optimizer = torch.optim.AdamW(list(classifier.parameters()) + list(encoder.parameters()), lr=1e-4, weight_decay=0.01) scheduler = get_linear_schedule_with_warmup( optimizer, num_warmup_steps= warm_up_steps, num_training_steps=total_steps ) logger.info("Starting training of MiniLM model with classification head...") low_loss = float('inf') patience_counter = 0 for epoch in range(epoch): encoder.train() classifier.train() losses = [] for i in range(0, len(train_set), batch_size): batch = train_set[i:i+batch_size] labels = torch.tensor(batch['labels']).to(device) texts = drop_tokens_batch(batch['text']) features = encoder.tokenize(texts) features = {k: v.to(device) for k, v in features.items()} embeddings = encoder(features)['sentence_embedding'] embeddings = torch.tensor(embeddings).to(device) logits = classifier(embeddings) loss = criterion(logits, labels) optimizer.zero_grad() loss.backward() optimizer.step() scheduler.step() if i % 100 == 0: logger.info("Done {} out of {} batches".format(i, len(train_set))) encoder.eval() classifier.eval() with torch.no_grad(): for i in range(0, len(eval_set), batch_size): batch = eval_set[i:i+batch_size] labels = torch.tensor(batch['labels']).to(device) embeddings = encoder.encode(batch['text']) embeddings = torch.tensor(embeddings).to(device) logits = classifier(embeddings) loss = criterion(logits, labels) losses.append(loss.item()) avg_loss = sum(losses) / len(losses) logger.info(f"Epoch {epoch+1} completed, Loss: {avg_loss:.4f}") if(avg_loss < low_loss): low_loss = avg_loss patience_counter = 0 encoder.save(model_save_path) logger.info(f"encoder saved at {model_save_path}.") else: patience_counter += 1 if(patience_counter >= 2): logger.info("Early stopping triggered.") break logger.info("MiniLM model trained with classification head.") return { "total_steps": total_steps, "warm_up_steps": warm_up_steps, "batch_size": batch_size, "epochs": epoch, "model_save_path": model_save_path } class MiniLMModel(BaseModel): """ MiniLM model implementation for efficient text embeddings. """ def __init__(self, language, path=None): """ Initialize the MiniLM model with configuration parameters. Args: language (str): Language for the model. path (str, optional): Path to load a pre-trained model. Defaults to None. If None, a new model is initialized. """ self.number_of_estimators = 300 self.learning_rate = 0.1 self.max_depth = 4 self.tree_method = 'hist' self.objective = 'binary:logistic' self.eval_metric = 'logloss' self.device = "cuda" if torch.cuda.is_available() else "cpu" self.params = { "number_of_estimators": self.number_of_estimators, "learning_rate": self.learning_rate, "max_depth": self.max_depth, "tree_method": self.tree_method, "objective": self.objective, "eval_metric": self.eval_metric } super().__init__(language, path) def setup_model(self): """ Initialize the MiniLM SentenceTransformer model. """ self.encoder = None self.model_path = "sentence-transformers/minilm.pt" xgb_classifier = XGBClassifier(n_estimators=self.number_of_estimators, eval_metric=self.eval_metric, objective=self.objective, learning_rate=self.learning_rate, max_depth=self.max_depth, tree_method=self.tree_method) self.classifier = MultiOutputClassifier(xgb_classifier) logger.info("MiniLM model initialized.") def train(self, X_train, y_train): """ Train the MiniLM model with a classification head. Args: X_train: Input training data. y_train: True labels for training data. """ if self.encoder is None and self.language == "java": if os.path.exists(self.model_path): logger.info(f"Loading existing MiniLM model from {self.model_path} for fine-tuning...") self.encoder = SentenceTransformer(self.model_path).to(self.device) else: logger.info(f"Fine-tuning MiniLM encoder using {self.language} training data...") parameters = finetune_miniLM(X_train, y_train,device=self.device, model_save_path=self.model_path) self.params.update(parameters) self.encoder = SentenceTransformer(parameters["model_save_path"]).to(self.device) if self.encoder is None: self.encoder = SentenceTransformer(self.model_path).to(self.device) y_train = np.array(y_train,dtype=np.float32) train_embeddings = self.encoder.encode(X_train) logger.info("Starting training of MiniLM model with Xgboost...") self.classifier.fit(train_embeddings, y_train) return { "n_estimators": self.number_of_estimators, "learning_rate": self.learning_rate, "max_depth": self.max_depth, "tree_method": self.tree_method, "objective": self.objective, "eval_metric": self.eval_metric } def evaluate(self, X_test, y_test) -> dict[str,any]: """ Evaluate the MiniLM model on test data. Args: X_test: Input test data. y_test: True labels for test data. """ y_test = np.array(y_test,dtype=np.float32) test_embeddings = self.encoder.encode(X_test) predictions = self.classifier.predict(test_embeddings) accuracy = accuracy_score(y_test, predictions) f1_micro = f1_score(y_test, predictions, average='micro') f1_macro = f1_score(y_test, predictions, average='macro') f1_weighted = f1_score(y_test, predictions, average='weighted') recall = recall_score(y_test, predictions, average='weighted') precision = precision_score(y_test, predictions, average='weighted') metrics = { "accuracy": accuracy, "f1_micro_score": f1_micro, "f1_macro_score": f1_macro, "f1_weighted_score": f1_weighted, "recall": recall, "precision": precision } return metrics def predict(self, X) -> ndarray: """ Make predictions using the trained MiniLM model. Args: X: Input data for prediction. Returns: Predictions made by the model. """ if self.encoder is None or self.classifier is None: raise ValueError("Model is not trained. Call train() or load() before prediction.") encodedText = self.encoder.encode(X) predictions = self.classifier.predict(encodedText) logger.info(f"Predictions: {predictions}.") return predictions def save(self, path, model_name): """ Save model and log to MLflow. Args: path (str): Path to save the model. model_name (str): Name to use when saving the model (without extension). """ if self.encoder is None and self.classifier is None: raise ValueError("Model is not trained. Cannot save uninitialized model.") complete_path = os.path.join(path, model_name) encoder_path = complete_path+f"_encoder_{self.language}" classifier_path = complete_path+f"_xgb_classifier_{self.language}.joblib" if os.path.exists(complete_path) and os.path.isdir(complete_path): shutil.rmtree(complete_path) self.encoder.save(encoder_path) joblib.dump(self.classifier, classifier_path) try: # Log to MLflow logger.info("Logging artifacts to MLflow...") mlflow.pyfunc.log_model( artifact_path=f"{model_name}_{self.language}", python_model=MiniLMClassifierWrapper(), artifacts={ "encoder_path": encoder_path, "classifier_path": classifier_path }, code_paths=["turing/modeling/models/MiniLMClassifierWrapper.py"] ) except Exception as e: logger.error(f"Failed to log model artifacts to MLflow: {e}")