Upload 2 files

inference.py

@@ -0,0 +1,21 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+
+# Path of fine-tuned model
+model_path = "./fine_tuned_model"
+
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = AutoModelForCausalLM.from_pretrained(model_path)
+
+# Create chatbot pipeline
+chatbot = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device=0 if torch.cuda.is_available() else -1  # Use GPU if available
+)
+
+# Example usage
+prompt = "Hello, can you tell me some fun facts about european legislation?"
+response = chatbot(prompt, max_length=100, do_sample=True, temperature=0.7)
+print(response[0]['generated_text'])

train.py

@@ -0,0 +1,85 @@
+from datasets import load_dataset
+from transformers import DataCollatorForLanguageModeling
+from transformers import Trainer, TrainingArguments
+import os
+import torch
+
+
+
+def main():
+      
+    local_rank = int(os.environ['LOCAL_RANK'])    
+    rank = int(os.environ['RANK'])    
+    world_size = int(os.environ['WORLD_SIZE']) 
+
+    torch.distributed.init_process_group("nccl")    
+    print(f"Local Rank = {local_rank}/{world_size}")    
+
+
+
+    # Load your JSONL file
+    dataset = load_dataset('json', data_files='../../data/m500_clean.jsonl', split='train')
+    
+    # Load a Model
+    from transformers import AutoTokenizer, AutoModelForCausalLM
+    
+    model_name = "FacebookAI/roberta-base"
+    
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    
+    # Set pad token if not set
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    
+    # Tokenize the dataset
+    def tokenize_function(examples):
+        return tokenizer(examples["text"], truncation=True, max_length=512)
+    
+    tokenized_dataset = dataset.map(tokenize_function, batched=True)
+    
+    # Split the dataset into training and validation sets
+    split_dataset = tokenized_dataset.train_test_split(test_size=0.1)
+    
+    # Data collator, pad the inputs to the maximum length in the batch
+    data_collator = DataCollatorForLanguageModeling(
+        tokenizer=tokenizer, mlm=False  # mlm=False: causal language modeling
+    )
+    
+    # Training
+    training_args = TrainingArguments(
+        output_dir="./results",
+        overwrite_output_dir=True,
+        num_train_epochs=3,
+        per_device_train_batch_size=4,
+        per_device_eval_batch_size=4,
+        dataloader_num_workers=8,  
+        eval_steps=500,
+        save_steps=1000,
+        warmup_steps=500,
+        prediction_loss_only=True,
+        logging_dir="./logs",
+        logging_steps=100,
+        learning_rate=5e-5,
+        fp16=True,  # true for GPU
+    )
+    
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=split_dataset["train"],
+        eval_dataset=split_dataset["test"],
+        data_collator=data_collator,
+    )
+    
+    # Start training
+    trainer.train()
+
+    torch.distributed.destroy_process_group()
+    
+    # Save the model and tokenizer
+    model.save_pretrained("./fine_tuned_model")
+    tokenizer.save_pretrained("./fine_tuned_model")
+
+if __name__ == "__main__":
+    main()