MindsDB 25.7.4.0__py3-none-any.whl → 25.8.3.0__py3-none-any.whl

mindsdb/integrations/handlers/huggingface_handler/finetune.py

@@ -1,223 +1,203 @@
-# import evaluate
-# import nltk
-# import numpy as np
-# from datasets import Dataset
-# from transformers import (
-#     AutoConfig,
-#     AutoModelForSeq2SeqLM,
-#     AutoModelForSequenceClassification,
-#     AutoTokenizer,
-#     DataCollatorForSeq2Seq,
-#     Seq2SeqTrainingArguments,
-#     Trainer,
-#     TrainingArguments,
-# )
-
-# # todo add support for question answering task
-# # todo add support for fill mask
-# # todo add support for text_generation (causal language model)
-# # todo add support for text_2_text generation
-
-
-# def _finetune_cls(df, args):
-#     df = df.rename(columns={args["target"]: "labels", args["input_column"]: "text"})
-#     tokenizer_from = args.get("using", {}).get("tokenizer_from", args["model_name"])
-#     tokenizer = AutoTokenizer.from_pretrained(tokenizer_from)
-#     dataset = Dataset.from_pandas(df)
-
-#     def _tokenize_text_cls_fn(examples):
-#         return tokenizer(examples["text"], padding="max_length", truncation=True)
-
-#     tokenized_datasets = dataset.map(_tokenize_text_cls_fn, batched=True)
-#     ds = tokenized_datasets.shuffle(seed=42).train_test_split(
-#         test_size=args.get("eval_size", 0.1)
-#     )
-#     train_ds = ds["train"]
-#     eval_ds = ds["test"]
-
-#     ft_args = args.get("using", {}).get("trainer_args", {})
-#     ft_args["output_dir"] = args["model_folder"]
-
-#     n_labels = len(args["labels_map"])
-#     # todo replace for prod
-#     assert (
-#         n_labels == df["labels"].nunique()
-#     ), f'Label mismatch! Ensure labels match what the model was originally trained on. Found {df["labels"].nunique()} classes, expected {n_labels}.'  # noqa
-#     # TODO: ideally check that labels are a subset of the original ones, too.
-#     config = AutoConfig.from_pretrained(args["model_name"])
-#     model = AutoModelForSequenceClassification.from_pretrained(
-#         args["model_name"], config=config
-#     )
-#     metric = evaluate.load("accuracy")
-#     training_args = TrainingArguments(**ft_args)
-
-#     def _compute_metrics(eval_pred):
-#         logits, labels = eval_pred
-#         predictions = np.argmax(logits, axis=-1)
-#         return metric.compute(predictions=predictions, references=labels)
-
-#     # generate trainer and finetune
-#     trainer = Trainer(
-#         model=model,
-#         args=training_args,
-#         train_dataset=train_ds,
-#         eval_dataset=eval_ds,
-#         compute_metrics=_compute_metrics,
-#     )
-
-#     return tokenizer, trainer
-
-
-# # TODO: merge with summarization?
-# def _finetune_translate(df, args):
-#     config = AutoConfig.from_pretrained(args["model_name"])
-#     df = df.rename(
-#         columns={args["target"]: "translation", args["input_column"]: "text"}
-#     )
-#     tokenizer_from = args.get("using", {}).get("tokenizer_from", args["model_name"])
-#     tokenizer = AutoTokenizer.from_pretrained(tokenizer_from)
-#     dataset = Dataset.from_pandas(df)
-
-#     def _tokenize_translate_fn(examples):
-#         source_lang = args["lang_input"]
-#         target_lang = args["lang_output"]
-#         max_target_length = config.task_specific_params["summarization"]["max_length"]
-#         prefix = f"translate {source_lang} to {target_lang}: "
-#         inputs = [prefix + ex for ex in examples["text"]]
-#         targets = [ex for ex in examples["translation"]]
-#         model_inputs = tokenizer(inputs, max_length=config.n_positions, truncation=True)
-
-#         # Setup the tokenizer for targets
-#         with tokenizer.as_target_tokenizer():
-#             labels = tokenizer(targets, max_length=max_target_length, truncation=True)
-
-#         model_inputs["labels"] = labels["input_ids"]
-#         return model_inputs
-
-#     tokenized_datasets = dataset.map(_tokenize_translate_fn, batched=True)
-#     ds = tokenized_datasets.shuffle(seed=42).train_test_split(
-#         test_size=args.get("eval_size", 0.1)
-#     )
-#     train_ds = ds["train"]
-#     eval_ds = ds["test"]
-#     ft_args = args.get("using", {}).get("trainer_args", {})
-#     ft_args["output_dir"] = args["model_folder"]
-#     ft_args["predict_with_generate"] = True
-
-#     model = AutoModelForSeq2SeqLM.from_pretrained(args["model_name"], config=config)
-#     model.resize_token_embeddings(len(tokenizer))
-#     training_args = Seq2SeqTrainingArguments(**ft_args)
-#     data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
-
-#     # generate trainer and finetune
-#     trainer = Trainer(
-#         model=model,
-#         args=training_args,
-#         train_dataset=train_ds,
-#         eval_dataset=eval_ds,
-#         data_collator=data_collator,
-#         # compute_metrics=_compute_metrics,
-#     )
-
-#     return tokenizer, trainer
-
-
-# def _finetune_summarization(df, args):
-#     df = df.rename(columns={args["target"]: "summary", args["input_column"]: "text"})
-#     tokenizer_from = args.get("using", {}).get("tokenizer_from", args["model_name"])
-#     tokenizer = AutoTokenizer.from_pretrained(tokenizer_from)
-#     dataset = Dataset.from_pandas(df)
-#     config = AutoConfig.from_pretrained(args["model_name"])
-
-#     def _tokenize_summarize_fn(examples):
-#         prefix = "summarize: " if "t5" in args["model_name"] else ""
-#         inputs = [prefix + doc for doc in examples["text"]]
-#         model_inputs = tokenizer(
-#             inputs,
-#             padding="max_length",
-#             truncation=True,
-#             max_length=config.max_position_embeddings,
-#             pad_to_max_length=True,
-#         )  # noqa
-#         labels = tokenizer(
-#             text_target=examples["summary"],
-#             max_length=config.max_position_embeddings,
-#             truncation=True,
-#         )  # noqa
-#         model_inputs["labels"] = labels["input_ids"]
-#         return model_inputs
-
-#     tokenized_datasets = dataset.map(_tokenize_summarize_fn, batched=True)
-#     ds = tokenized_datasets.shuffle(seed=42).train_test_split(
-#         test_size=args.get("eval_size", 0.1)
-#     )
-#     train_ds = ds["train"]
-#     eval_ds = ds["test"]
-
-#     ft_args = args.get("using", {}).get("trainer_args", {})
-#     ft_args["output_dir"] = args["model_folder"]
-#     ft_args["predict_with_generate"] = True
-
-#     model = AutoModelForSeq2SeqLM.from_pretrained(args["model_name"], config=config)
-#     metric = evaluate.load("rouge")
-#     training_args = Seq2SeqTrainingArguments(**ft_args)
-#     data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
-
-#     def _compute_metrics(eval_pred):
-#         # ref: github.com/huggingface/notebooks/blob/main/examples/summarization.ipynb
-#         predictions, labels = eval_pred
-#         decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
-#         decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
-
-#         # Rogue expects a newline after each sentence
-#         decoded_preds = [
-#             "\n".join(nltk.sent_tokenize(pred.strip())) for pred in decoded_preds
-#         ]
-#         decoded_labels = [
-#             "\n".join(nltk.sent_tokenize(label.strip())) for label in decoded_labels
-#         ]
-
-#         result = metric.compute(
-#             predictions=decoded_preds,
-#             references=decoded_labels,
-#             use_stemmer=True,
-#             use_aggregator=True,
-#         )
-#         result = {key: value * 100 for key, value in result.items()}
-#         prediction_lens = [
-#             np.count_nonzero(pred != tokenizer.pad_token_id) for pred in predictions
-#         ]
-#         result["gen_len"] = np.mean(prediction_lens)  # todo: remove?
-#         return {k: round(v, 4) for k, v in result.items()}
-
-#     # generate trainer and finetune
-#     trainer = Trainer(
-#         model=model,
-#         args=training_args,
-#         train_dataset=train_ds,
-#         eval_dataset=eval_ds,
-#         data_collator=data_collator,
-#         compute_metrics=_compute_metrics,
-#     )
-
-#     return tokenizer, trainer
-
-
-# def _finetune_fill_mask(df, args):
-#     raise NotImplementedError("Finetuning fill-mask models is not yet supported.")
-
-
-# def _finetune_text_generation(df, args):
-#     raise NotImplementedError("Finetuning text-generation models is not yet supported.")
-
-
-# def _finetune_question_answering(df, args):
-#     raise NotImplementedError(
-#         "Finetuning question-answering models is not yet supported."
-#     )
-
-
-# def _finetune_text_2_text_generation(df, args):
-#     raise NotImplementedError(
-#         "Finetuning text-2-text generation models is not yet supported."
-#     )
+import evaluate
+import nltk
+import numpy as np
+from datasets import Dataset
+from transformers import (
+    AutoConfig,
+    AutoModelForSeq2SeqLM,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    DataCollatorForSeq2Seq,
+    Seq2SeqTrainingArguments,
+    Trainer,
+    TrainingArguments,
+)
+
+# todo add support for question answering task
+# todo add support for fill mask
+# todo add support for text_generation (causal language model)
+# todo add support for text_2_text generation
+
+
+def _finetune_cls(df, args):
+    df = df.rename(columns={args["target"]: "labels", args["input_column"]: "text"})
+    tokenizer_from = args.get("using", {}).get("tokenizer_from", args["model_name"])
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_from)
+    dataset = Dataset.from_pandas(df)
+
+    def _tokenize_text_cls_fn(examples):
+        return tokenizer(examples["text"], padding="max_length", truncation=True)
+
+    tokenized_datasets = dataset.map(_tokenize_text_cls_fn, batched=True)
+    ds = tokenized_datasets.shuffle(seed=42).train_test_split(test_size=args.get("eval_size", 0.1))
+    train_ds = ds["train"]
+    eval_ds = ds["test"]
+
+    ft_args = args.get("using", {}).get("trainer_args", {})
+    ft_args["output_dir"] = args["model_folder"]
+
+    n_labels = len(args["labels_map"])
+    # todo replace for prod
+    assert n_labels == df["labels"].nunique(), (
+        f"Label mismatch! Ensure labels match what the model was originally trained on. Found {df['labels'].nunique()} classes, expected {n_labels}."
+    )  # noqa
+    # TODO: ideally check that labels are a subset of the original ones, too.
+    config = AutoConfig.from_pretrained(args["model_name"])
+    model = AutoModelForSequenceClassification.from_pretrained(args["model_name"], config=config)
+    metric = evaluate.load("accuracy")
+    training_args = TrainingArguments(**ft_args)
+
+    def _compute_metrics(eval_pred):
+        logits, labels = eval_pred
+        predictions = np.argmax(logits, axis=-1)
+        return metric.compute(predictions=predictions, references=labels)
+
+    # generate trainer and finetune
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_ds,
+        eval_dataset=eval_ds,
+        compute_metrics=_compute_metrics,
+    )
+
+    return tokenizer, trainer
+
+
+# TODO: merge with summarization?
+def _finetune_translate(df, args):
+    config = AutoConfig.from_pretrained(args["model_name"])
+    df = df.rename(columns={args["target"]: "translation", args["input_column"]: "text"})
+    tokenizer_from = args.get("using", {}).get("tokenizer_from", args["model_name"])
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_from)
+    dataset = Dataset.from_pandas(df)
+
+    def _tokenize_translate_fn(examples):
+        source_lang = args["lang_input"]
+        target_lang = args["lang_output"]
+        max_target_length = config.task_specific_params["summarization"]["max_length"]
+        prefix = f"translate {source_lang} to {target_lang}: "
+        inputs = [prefix + ex for ex in examples["text"]]
+        targets = [ex for ex in examples["translation"]]
+        model_inputs = tokenizer(inputs, max_length=config.n_positions, truncation=True)
+
+        # Setup the tokenizer for targets
+        with tokenizer.as_target_tokenizer():
+            labels = tokenizer(targets, max_length=max_target_length, truncation=True)
+
+        model_inputs["labels"] = labels["input_ids"]
+        return model_inputs
+
+    tokenized_datasets = dataset.map(_tokenize_translate_fn, batched=True)
+    ds = tokenized_datasets.shuffle(seed=42).train_test_split(test_size=args.get("eval_size", 0.1))
+    train_ds = ds["train"]
+    eval_ds = ds["test"]
+    ft_args = args.get("using", {}).get("trainer_args", {})
+    ft_args["output_dir"] = args["model_folder"]
+    ft_args["predict_with_generate"] = True
+
+    model = AutoModelForSeq2SeqLM.from_pretrained(args["model_name"], config=config)
+    model.resize_token_embeddings(len(tokenizer))
+    training_args = Seq2SeqTrainingArguments(**ft_args)
+    data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
+
+    # generate trainer and finetune
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_ds,
+        eval_dataset=eval_ds,
+        data_collator=data_collator,
+        # compute_metrics=_compute_metrics,
+    )
+
+    return tokenizer, trainer
+
+
+def _finetune_summarization(df, args):
+    df = df.rename(columns={args["target"]: "summary", args["input_column"]: "text"})
+    tokenizer_from = args.get("using", {}).get("tokenizer_from", args["model_name"])
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_from)
+    dataset = Dataset.from_pandas(df)
+    config = AutoConfig.from_pretrained(args["model_name"])
+
+    def _tokenize_summarize_fn(examples):
+        prefix = "summarize: " if "t5" in args["model_name"] else ""
+        inputs = [prefix + doc for doc in examples["text"]]
+        model_inputs = tokenizer(
+            inputs,
+            padding="max_length",
+            truncation=True,
+            max_length=config.max_position_embeddings,
+            pad_to_max_length=True,
+        )  # noqa
+        labels = tokenizer(
+            text_target=examples["summary"],
+            max_length=config.max_position_embeddings,
+            truncation=True,
+        )  # noqa
+        model_inputs["labels"] = labels["input_ids"]
+        return model_inputs
+
+    tokenized_datasets = dataset.map(_tokenize_summarize_fn, batched=True)
+    ds = tokenized_datasets.shuffle(seed=42).train_test_split(test_size=args.get("eval_size", 0.1))
+    train_ds = ds["train"]
+    eval_ds = ds["test"]
+
+    ft_args = args.get("using", {}).get("trainer_args", {})
+    ft_args["output_dir"] = args["model_folder"]
+    ft_args["predict_with_generate"] = True
+
+    model = AutoModelForSeq2SeqLM.from_pretrained(args["model_name"], config=config)
+    metric = evaluate.load("rouge")
+    training_args = Seq2SeqTrainingArguments(**ft_args)
+    data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
+
+    def _compute_metrics(eval_pred):
+        # ref: github.com/huggingface/notebooks/blob/main/examples/summarization.ipynb
+        predictions, labels = eval_pred
+        decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
+        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
+
+        # Rogue expects a newline after each sentence
+        decoded_preds = ["\n".join(nltk.sent_tokenize(pred.strip())) for pred in decoded_preds]
+        decoded_labels = ["\n".join(nltk.sent_tokenize(label.strip())) for label in decoded_labels]
+
+        result = metric.compute(
+            predictions=decoded_preds,
+            references=decoded_labels,
+            use_stemmer=True,
+            use_aggregator=True,
+        )
+        result = {key: value * 100 for key, value in result.items()}
+        prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in predictions]
+        result["gen_len"] = np.mean(prediction_lens)  # todo: remove?
+        return {k: round(v, 4) for k, v in result.items()}
+
+    # generate trainer and finetune
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_ds,
+        eval_dataset=eval_ds,
+        data_collator=data_collator,
+        compute_metrics=_compute_metrics,
+    )
+
+    return tokenizer, trainer
+
+
+def _finetune_fill_mask(df, args):
+    raise NotImplementedError("Finetuning fill-mask models is not yet supported.")
+
+
+def _finetune_text_generation(df, args):
+    raise NotImplementedError("Finetuning text-generation models is not yet supported.")
+
+
+def _finetune_question_answering(df, args):
+    raise NotImplementedError("Finetuning question-answering models is not yet supported.")
+
+
+def _finetune_text_2_text_generation(df, args):
+    raise NotImplementedError("Finetuning text-2-text generation models is not yet supported.")