Upload gpt2.py with huggingface_hub

gpt2.py

@@ -0,0 +1,94 @@
+
+from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, GPT2Config, GPT2LMHeadModel, pipeline
+import torch 
+torch.set_default_dtype(torch.float32) 
+
+class SLMGPT2:
+    def __init__(self, tokenizer_id, model_id):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.dataset = load_dataset("tiny_shakespeare", trust_remote_code=True)
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        
+        self.tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+        self.distilled_model = AutoModelForCausalLM.from_pretrained(model_id)
+        self.vocab_size=self.tokenizer.vocab_size
+        self.n_positions=128
+        self.n_ctx=128
+        self.n_embd=256
+        self.n_layer=4
+        self.n_head=4
+        print(f"device set: {self.device}")
+        self.config = GPT2Config(
+                        vocab_size=self.vocab_size,
+                        n_positions=self.n_positions,
+                        n_ctx=self.n_ctx,
+                        n_embd=self.n_embd,
+                        n_layer=self.n_layer,
+                        n_head=self.n_head
+                    )
+        self.config.pad_token_id = self.tokenizer.pad_token_id
+        print(f"Vocab size: {self.config.vocab_size}, type: {type(self.config.vocab_size)}")
+        print(f"Embedding dim: {self.config.n_embd}, type: {type(self.config.n_embd)}")
+        self.model = GPT2LMHeadModel(self.config).to(device=self.device)
+        print(f"device set: {self.device}")
+
+    def tokenize(self, example):
+        tokenized = self.tokenizer(
+            example["text"],
+            truncation=True,
+            padding="max_length",
+            max_length=128
+        )
+        tokenized["labels"] = tokenized["input_ids"].copy()
+        
+        # Mask padding tokens in the labels
+        pad_token_id = self.tokenizer.pad_token_id
+        tokenized["labels"] = [
+            (label if label != pad_token_id else -100)
+            for label in tokenized["labels"]
+        ]
+        return tokenized
+        
+    def train(self):
+        tokenized_dataset = self.dataset.map(self.tokenize, batched=True)
+        training_args = TrainingArguments(
+            output_dir="./results",
+            num_train_epochs=500,
+            per_device_train_batch_size=4,
+            per_device_eval_batch_size=4,
+            logging_steps=10,
+            save_steps=500,
+            learning_rate=5e-4,
+            weight_decay=0.01,
+            save_total_limit=1,
+            logging_dir="none"
+        )
+
+        trainer = Trainer(
+            model=self.model,
+            args=training_args,
+            train_dataset=tokenized_dataset["train"]
+        )
+        trainer.train()
+        print("training completed!!")
+        self.tokenizer.save_pretrained("./saved_model")
+        self.model.save_pretrained("./saved_model")
+        print("Model and tokenizer saved to ./saved_model")
+
+    def inference(self):
+        model_path = "./saved_model"
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
+        model = AutoModelForCausalLM.from_pretrained(model_path)
+        text_gen = pipeline("text-generation", model=model, tokenizer=tokenizer)
+        result = text_gen("The universe is", max_length=128)
+        print(result[0]["generated_text"])
+
+if __name__ == "__main__":
+    model_id = "gpt2"
+    tokenizer_id = "distilgpt2"
+    sml = SLMGPT2(model_id=model_id, tokenizer_id=tokenizer_id)
+    sml.train()
+    sml.inference()