Update train.py

train.py

@@ -2,11 +2,6 @@ from transformers import MBartForSequenceClassification, MBart50Tokenizer, Train
 from datasets import Dataset
 
 
-# Load the model and tokenizer
-model_name = "LocalDoc/mbart_large_qa_azerbaijan"  # Replace with your model name if different
-tokenizer = MBart50Tokenizer.from_pretrained(model_name)
-model = MBartForSequenceClassification.from_pretrained(model_name)
-chunk_size = 512
 
 # Prepare the dataset (simplified)
 def prepare_text_dataset(data):
@@ -29,34 +24,40 @@ def prepare_text_dataset(data):
 
     return formatted_dataset
 
-
-# Load the plain text (replace with your actual loading logic)
-with open("constitution.txt", "r", encoding="utf-8") as f:
-  constitution_text = f.read()
-
-# Prepare the dataset
-train_dataset = prepare_text_dataset(constitution_text)
-
-# Define training arguments
-training_args = TrainingArguments(
-    output_dir="./results",  # Adjust output directory
-    overwrite_output_dir=True,
-    num_train_epochs=3,  # Adjust training epochs
-    per_device_train_batch_size=1,  # Adjust batch size based on your GPU memory
-    save_steps=500,
-    save_total_limit=2,
-)
-
-# Create the Trainer
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    train_dataset=train_dataset,
-)
-
-# Start training
-trainer.train()
-
-# Save the fine-tuned model
-model.save_pretrained("./fine-tuned_model")
-tokenizer.save_pretrained("./fine-tuned_model")
+def init():
+    # Load the model and tokenizer
+    model_name = "LocalDoc/mbart_large_qa_azerbaijan"  # Replace with your model name if different
+    tokenizer = MBart50Tokenizer.from_pretrained(model_name)
+    model = MBartForSequenceClassification.from_pretrained(model_name)
+    chunk_size = 512
+
+    # Load the plain text (replace with your actual loading logic)
+    with open("constitution.txt", "r", encoding="utf-8") as f:
+      constitution_text = f.read()
+    
+    # Prepare the dataset
+    train_dataset = prepare_text_dataset(constitution_text)
+    
+    # Define training arguments
+    training_args = TrainingArguments(
+        output_dir="./results",  # Adjust output directory
+        overwrite_output_dir=True,
+        num_train_epochs=3,  # Adjust training epochs
+        per_device_train_batch_size=1,  # Adjust batch size based on your GPU memory
+        save_steps=500,
+        save_total_limit=2,
+    )
+    
+    # Create the Trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+    )
+    
+    # Start training
+    trainer.train()
+    
+    # Save the fine-tuned model
+    model.save_pretrained("./fine-tuned_model")
+    tokenizer.save_pretrained("./fine-tuned_model")