Create huggingface_helper.py

huggingface_helper.py

@@ -0,0 +1,51 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
+from datasets import load_dataset
+import os
+
+class HuggingFaceHelper:
+    def __init__(self, model_path="./merged_model", dataset_path=None):
+        self.model_path = model_path
+        self.dataset_path = dataset_path
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model = AutoModelForCausalLM.from_pretrained(model_path).to(self.device)
+
+    def load_dataset(self):
+        if self.dataset_path:
+            dataset = load_dataset("json", data_files=self.dataset_path, split="train")
+            return dataset.map(self.tokenize_function, batched=True)
+        raise ValueError("Dataset path not provided.")
+
+    def tokenize_function(self, examples):
+        return self.tokenizer(examples["messages"], truncation=True, padding="max_length", max_length=512)
+
+    def fine_tune(self, output_dir="./fine_tuned_model", epochs=3, batch_size=4):
+        dataset = self.load_dataset()
+
+        training_args = TrainingArguments(
+            output_dir=output_dir,
+            evaluation_strategy="epoch",
+            save_strategy="epoch",
+            per_device_train_batch_size=batch_size,
+            num_train_epochs=epochs,
+            weight_decay=0.01,
+            push_to_hub=True,
+            hub_model_id="Raiff1982/codriao-finetuned"
+        )
+
+        trainer = Trainer(
+            model=self.model,
+            args=training_args,
+            train_dataset=dataset,
+            tokenizer=self.tokenizer,
+        )
+
+        trainer.train()
+        self.save_model(output_dir)
+
+    def save_model(self, output_dir):
+        self.model.save_pretrained(output_dir)
+        self.tokenizer.save_pretrained(output_dir)
+        print(f"✅ Model saved to {output_dir} and uploaded to Hugging Face Hub.")