Create train_model.py

train_model.py

@@ -0,0 +1,149 @@
+# training_space/train_model.py (Training Script)
+import argparse
+from transformers import (
+    GPT2Config, GPT2LMHeadModel, 
+    BertConfig, BertForSequenceClassification,
+    Trainer, TrainingArguments, AutoTokenizer,
+    DataCollatorForLanguageModeling, DataCollatorWithPadding
+)
+from datasets import load_dataset, Dataset
+import torch
+import os
+from huggingface_hub import HfApi, HfFolder
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--task", type=str, required=True, help="Task type: generation or classification")
+    parser.add_argument("--model_name", type=str, required=True, help="Name of the model")
+    parser.add_argument("--dataset", type=str, required=True, help="Path to the dataset")
+    parser.add_argument("--num_layers", type=int, default=12)
+    parser.add_argument("--attention_heads", type=int, default=1)
+    parser.add_argument("--hidden_size", type=int, default=64)
+    parser.add_argument("--vocab_size", type=int, default=30000)
+    parser.add_argument("--sequence_length", type=int, default=512)
+    args = parser.parse_args()
+    
+    # Define output directory
+    output_dir = f"./models/{args.model_name}"
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Initialize Hugging Face API
+    api = HfApi()
+    hf_token = HfFolder.get_token()
+    
+    # Initialize tokenizer (adjust based on task)
+    if args.task == "generation":
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    elif args.task == "classification":
+        tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+    else:
+        raise ValueError("Unsupported task type")
+    
+    # Load and prepare dataset
+    if args.task == "generation":
+        dataset = load_dataset('text', data_files={'train': args.dataset})
+        def tokenize_function(examples):
+            return tokenizer(examples['text'], truncation=True, max_length=args.sequence_length)
+    elif args.task == "classification":
+        # For classification, assume the dataset is a simple text file with "text\tlabel" per line
+        with open(args.dataset, "r", encoding="utf-8") as f:
+            lines = f.readlines()
+        texts = []
+        labels = []
+        for line in lines:
+            parts = line.strip().split("\t")
+            if len(parts) == 2:
+                texts.append(parts[0])
+                labels.append(int(parts[1]))
+        dataset = Dataset.from_dict({"text": texts, "label": labels})
+        def tokenize_function(examples):
+            return tokenizer(examples['text'], truncation=True, max_length=args.sequence_length)
+    else:
+        raise ValueError("Unsupported task type")
+    
+    tokenized_datasets = dataset.map(tokenize_function, batched=True)
+    
+    if args.task == "generation":
+        data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+    elif args.task == "classification":
+        data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+    
+    # Initialize model based on task
+    if args.task == "generation":
+        config = GPT2Config(
+            vocab_size=args.vocab_size,
+            n_positions=args.sequence_length,
+            n_ctx=args.sequence_length,
+            n_embd=args.hidden_size,
+            num_hidden_layers=args.num_layers,
+            num_attention_heads=args.attention_heads,
+            intermediate_size=4 * args.hidden_size,
+            hidden_act='gelu',
+            use_cache=True
+        )
+        model = GPT2LMHeadModel(config)
+    elif args.task == "classification":
+        config = BertConfig(
+            vocab_size=args.vocab_size,
+            max_position_embeddings=args.sequence_length,
+            hidden_size=args.hidden_size,
+            num_hidden_layers=args.num_layers,
+            num_attention_heads=args.attention_heads,
+            intermediate_size=4 * args.hidden_size,
+            hidden_act='gelu',
+            num_labels=2  # Adjust based on your classification task
+        )
+        model = BertForSequenceClassification(config)
+    else:
+        raise ValueError("Unsupported task type")
+    
+    # Define training arguments
+    if args.task == "generation":
+        training_args = TrainingArguments(
+            output_dir=output_dir,
+            num_train_epochs=3,
+            per_device_train_batch_size=8,
+            save_steps=5000,
+            save_total_limit=2,
+            logging_steps=500,
+            learning_rate=5e-4,
+            remove_unused_columns=False
+        )
+    elif args.task == "classification":
+        training_args = TrainingArguments(
+            output_dir=output_dir,
+            num_train_epochs=3,
+            per_device_train_batch_size=16,
+            evaluation_strategy="epoch",
+            save_steps=5000,
+            save_total_limit=2,
+            logging_steps=500,
+            learning_rate=5e-5,
+            remove_unused_columns=False
+        )
+    
+    # Initialize Trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_datasets['train'],
+        data_collator=data_collator,
+    )
+    
+    # Start training
+    trainer.train()
+    
+    # Save the final model
+    trainer.save_model(output_dir)
+    tokenizer.save_pretrained(output_dir)
+    
+    # Push to Hugging Face Hub
+    model_repo = f"your-username/{args.model_name}"
+    api.create_repo(repo_id=model_repo, private=False, token=hf_token)
+    model.push_to_hub(model_repo, use_auth_token=hf_token)
+    tokenizer.push_to_hub(model_repo, use_auth_token=hf_token)
+    
+    print(f"Model '{args.model_name}' trained and pushed to Hugging Face Hub at '{model_repo}'.")
+
+if __name__ == "__main__":
+    main()