Update README.md

README.md

@@ -41,14 +41,22 @@ Loading the Model and Tokenizer
 ```python
 import torch
 from transformers import GPT2TokenizerFast
+import os
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from transformers import AutoTokenizer
+from datasets import load_dataset, concatenate_datasets, Dataset
+from tqdm import tqdm
+import math
+from huggingface_hub import hf_hub_download
 from huggingface_hub import snapshot_download
 
 # Download the model from Hugging Face Hub
 model_path = snapshot_download(repo_id="liminerity/tiny-epstein-100m")
 
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained(model_path)
-
+tokenizer = GPT2TokenizerFast.from_pretrained('gpt2')
+tokenizer.pad_token = tokenizer.eos_token
 # ------------------------------------------------------------------------------
 # Configuration (scaled to ~150M for L4 GPU)
 # ------------------------------------------------------------------------------
@@ -263,27 +271,23 @@ class TinyAya(nn.Module):
             next_token = torch.multinomial(probs, num_samples=1)
             input_ids = torch.cat([input_ids, next_token], dim=-1)
         return input_ids
-tokenizer = GPT2TokenizerFast.from_pretrained(repo_id)
-tokenizer.pad_token = tokenizer.eos_token
+
 model = TinyAya(ModelConfig())
 state_dict = torch.load(os.path.join(model_path, "pytorch_model.bin"), map_location="cpu")
 model.load_state_dict(state_dict)
 model.eval()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 ```
 
 Text Generation Example
 
 ```python
-prompt = "The Epstein files reveal"
-inputs = tokenizer(prompt, return_tensors="pt")
+prompt = """Was Jeffrey a good guy?"""
+input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
 with torch.no_grad():
-    outputs = model.generate(
-        inputs.input_ids,
-        max_new_tokens=50,
-        temperature=0.8,
-        do_sample=True
-    )
-print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+    output = model.generate(input_ids, max_new_tokens=50, temperature=0.8)
+print("Generated text:")
+print(tokenizer.decode(output[0]))
 ```
 
 Training Details