Added application file

app.py

@@ -0,0 +1,184 @@
+import gradio as gr
+import torch
+from src.model import Config,GPT
+from src.inference import GPTInfer
+import tiktoken
+import torch
+import os
+from huggingface_hub import hf_hub_download
+
+os.environ['GRADIO_DEFAULT_CONCURRENCY_LIMIT']="1"
+
+device = 'cpu'
+if torch.cuda.is_available():
+    device = 'cuda'
+elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+    device = 'mps'   
+print(f'using device: {device}')
+model_path = hf_hub_download(
+    repo_id="VJyzCELERY/GPT2-GutenbergStoryGenerator",  
+    filename="GPT2-GutenbergStoryGenerator.pt"                 
+)
+checkpoint = torch.load(model_path, weights_only=False)
+model = GPT(config=checkpoint['config'])
+model.load_state_dict(checkpoint['model'])
+model = model.to(device)
+token_encoder = tiktoken.get_encoding('gpt2')
+generator = GPTInfer(model, token_encoder, device)
+
+def generate_story(
+    prompt,
+    max_new_tokens=50,
+    seed=42,
+    temperature=0.8,
+    top_k=None,
+    top_p=0.9,
+    repetition_penalty=1.2,
+    frequency_penalty=0.6,
+    no_repeat_ngram_size=3,
+    longer_story=True,
+    context_window=512
+):
+    if not prompt.strip():
+        return prompt, gr.update()
+    
+    if top_k <= 0:
+        top_k = None
+    
+    output_text = prompt
+    last_piece = ""
+    # print(f'{prompt}',end='',flush=True)
+    yield gr.update(value=output_text,interactive=False), gr.update(interactive=False)  
+    for piece in generator.generate(
+        prompt,
+        max_new_tokens=max_new_tokens,
+        seed=seed,
+        temperature=temperature,
+        top_k=top_k,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        frequency_penalty=frequency_penalty,
+        no_repeat_ngram_size=no_repeat_ngram_size,
+        longer_story=longer_story,
+        context_window=context_window
+    ):
+        if piece == last_piece:
+            continue
+        last_piece = piece
+        output_text += piece
+        # print(f'{piece}',end='',flush=True)
+        yield output_text, gr.update(interactive=False)  
+    
+    yield gr.update(value=output_text,interactive=True), gr.update(interactive=True)
+
+with gr.Blocks(title="Story Generator") as demo:
+
+    gr.Markdown("# ✨ Story Generator ✨")
+    gr.Markdown(
+        "Ketik prompt atau cerita awal di bawah ini. "
+        "Tekan **Generate** untuk melanjutkan cerita. "
+        "Anda dapat mengedit hasil cerita dan generate lagi untuk melanjutkan."
+    )
+
+    story_box = gr.Textbox(
+        label="Story / Prompt",
+        lines=15,
+        placeholder="Tulis prompt atau awal cerita di sini...",
+    )
+
+    generate_btn = gr.Button("Generate Story", variant="primary")
+
+    with gr.Accordion("Generation Settings", open=False):
+        context_window = gr.Slider(
+            minimum=128,
+            maximum=2048,
+            value=512,
+            step=64,
+            label="Context Window (tokens to use from end of text)",
+            info="Limits how much previous text is used. Lower = faster but less context."
+        )
+        
+        max_new_tokens = gr.Slider(
+            minimum=20,
+            maximum=2048,
+            value=1024,
+            step=10,
+            label="Max New Tokens"
+        )    
+
+        seed = gr.Number(
+        value=42, 
+        label="Seed"
+        )
+
+        temperature = gr.Slider(
+            minimum=0.1,
+            maximum=1.0,
+            value=0.8,
+            step=0.05,
+            label="Temperature"
+        )
+
+        top_k = gr.Slider(
+            minimum=0, 
+            maximum=200, 
+            value=0, 
+            step=1, 
+            label="Top-K (0 = disabled)"
+        )
+
+        top_p = gr.Slider(
+            minimum=0.0, 
+            maximum=1.0, 
+            value=0.9, 
+            step=0.01, 
+            label="Top-P"
+        )
+
+        repetition_penalty = gr.Slider(
+            minimum=1.0,
+            maximum=2.0,
+            value=1.2,
+            step=0.05,
+            label="Repetition Penalty"
+        )
+
+        frequency_penalty = gr.Slider(
+            minimum=0.0,
+            maximum=1.0,
+            value=0.6,
+            step=0.05,
+            label="Frequency Penalty"
+        )
+
+        no_repeat = gr.Slider(
+            minimum=1, 
+            maximum=10, 
+            value=3, 
+            step=1, 
+            label="No-Repeat N-gram Size"
+        )
+    
+
+    generate_btn.click(
+        fn=generate_story,
+        inputs=[
+            story_box, 
+            max_new_tokens, 
+            seed, 
+            temperature, 
+            top_k, 
+            top_p, 
+            repetition_penalty, 
+            frequency_penalty, 
+            no_repeat,
+            gr.Checkbox(value=True, visible=False),  
+            context_window  
+        ],
+        outputs=[story_box, generate_btn]
+    )
+
+#Run App
+if __name__ == "__main__":
+    
+    demo.launch(share=False)

src/inference.py

@@ -0,0 +1,226 @@
+import torch
+import torch.nn.functional as F
+
+def concat(prev, new):
+    if prev and prev[-1].isalnum() and new and new[0].isalnum():
+        return prev + " " + new
+    return prev + new
+
+class GPTInfer:
+    def __init__(self, model, token_encoder, device):
+        self.model = model
+        self.token_encoder = token_encoder
+        self.device = device
+        self.device_type = 'cuda' if device.startswith('cuda') else 'cpu'
+    
+    def get_token_length(self, text):
+        return len(self.token_encoder.encode(text, allowed_special={"<|endoftext|>"}))
+    
+    def apply_frequency_penalty_and_blocking(
+        self,
+        logits,                 
+        gen_tokens,            
+        frequency_penalty=0.5,  
+        no_repeat_ngram_size=3, 
+    ):
+        logits = logits.clone().float() 
+
+        if frequency_penalty and frequency_penalty > 0.0:
+            counts = {}
+            for t in gen_tokens[0].tolist():
+                counts[t] = counts.get(t, 0) + 1
+            if counts:
+                vocab_size = logits.shape[-1]
+                penalty = torch.zeros(vocab_size, dtype=logits.dtype, device=logits.device)
+                for tok, c in counts.items():
+                    if 0 <= tok < vocab_size:
+                        penalty[tok] = float(c) * float(frequency_penalty)
+                logits = logits - penalty.unsqueeze(0)
+
+        if no_repeat_ngram_size and no_repeat_ngram_size > 0:
+            n = no_repeat_ngram_size
+            cur = gen_tokens[0].tolist()
+            if len(cur) >= n - 1:
+                banned_next = set()
+                for i in range(len(cur) - (n - 1)):
+                    ngram = tuple(cur[i:i + n])
+                    prefix = tuple(ngram[:-1])
+                    banned_next.add(ngram[-1])
+                last_prefix = tuple(cur[-(n - 1):]) if n > 1 else tuple()
+                for i in range(len(cur) - (n - 1)):
+                    if tuple(cur[i:i + (n - 1)]) == last_prefix and i + (n - 1) < len(cur):
+                        banned_token = cur[i + (n - 1)]
+                        if 0 <= banned_token < logits.shape[-1]:
+                            logits[0, banned_token] = -1e9
+
+        return logits
+    
+    def sample_next_token(
+        self,
+        logits,                 
+        gen_tokens,             
+        seed_rng,              
+        temperature=0.8,
+        top_k=None,
+        top_p=0.9,
+        repetition_penalty=1.2,
+        frequency_penalty=0.5,
+        no_repeat_ngram_size=3,
+        recent_tokens_window=200,
+    ):
+        logits = logits.clone().float()
+
+        recent = gen_tokens[0, -recent_tokens_window:].tolist()
+        if repetition_penalty is not None and repetition_penalty != 1.0:
+            for t in set(recent):
+                if 0 <= t < logits.shape[-1]:
+                    logits[0, t] /= float(repetition_penalty)
+
+        logits = self.apply_frequency_penalty_and_blocking(
+            logits,
+            gen_tokens,
+            frequency_penalty=frequency_penalty,
+            no_repeat_ngram_size=no_repeat_ngram_size,
+        )
+
+        if temperature is not None and temperature != 1.0:
+            logits = logits / float(temperature)
+
+        sorted_logits, sorted_idx = torch.sort(logits, descending=True)
+        sorted_probs = F.softmax(sorted_logits, dim=-1)
+
+        if top_k is not None:
+            k = min(int(top_k), sorted_logits.shape[-1])
+            sorted_logits = sorted_logits[:, :k]
+            sorted_idx = sorted_idx[:, :k]
+            sorted_probs = sorted_probs[:, :k]
+
+        if top_p is not None and 0.0 < top_p < 1.0:
+            cum_probs = torch.cumsum(sorted_probs, dim=-1)
+            mask = cum_probs <= top_p
+            if not mask.any():
+                mask[0, 0] = True
+            keep_count = int(mask.sum(dim=-1).item())
+            sorted_probs = sorted_probs[:, :keep_count]
+            sorted_idx = sorted_idx[:, :keep_count]
+
+        sorted_probs = sorted_probs / (sorted_probs.sum(dim=-1, keepdim=True) + 1e-12)
+        next_index_in_sorted = torch.multinomial(sorted_probs, 1, generator=seed_rng)  
+        next_tok = sorted_idx.gather(-1, next_index_in_sorted)
+
+        return int(next_tok.item())
+    
+    def generate(
+        self,
+        prompt,
+        max_new_tokens=50,
+        seed=42,
+        longer_story=True,
+        temperature=0.8,
+        top_k=None,
+        top_p=0.9,
+        repetition_penalty=1.2,
+        frequency_penalty=0.5,
+        no_repeat_ngram_size=3,
+        context_window=None,
+        stream=True,
+    ):
+        self.model.eval()
+
+        tokens = self.token_encoder.encode(prompt)
+        if context_window is not None and len(tokens) > context_window:
+            tokens = tokens[-context_window:]
+        
+        tokens = torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(self.device)
+        gen_tokens = tokens.clone()
+        
+        if seed is not None:
+            sample_rng = torch.Generator(device=self.device).manual_seed(seed)
+        else:
+            sample_rng = torch.Generator(device=self.device)
+        
+        eos_id = self.token_encoder.encode("<|endoftext|>", allowed_special={"<|endoftext|>"})[0]
+        context_len = self.model.config.context_length
+        new_tokens_generated = 0
+        HARD_MAX_TOTAL = context_len + max_new_tokens + 10
+
+        while new_tokens_generated < max_new_tokens and gen_tokens.shape[1] < HARD_MAX_TOTAL:
+            if gen_tokens.shape[1] > context_len:
+                idx_cond = gen_tokens[:, -context_len:]
+            else:
+                idx_cond = gen_tokens
+
+            with torch.no_grad():
+                try:
+                    with torch.autocast(device_type=self.device_type, dtype=torch.bfloat16):
+                        logits, _ = self.model(idx_cond)
+                except Exception:
+                    logits, _ = self.model(idx_cond)
+
+            next_logits = logits[:, -1:, :].squeeze(1)
+
+            if longer_story and new_tokens_generated < 5:
+                next_logits[0, eos_id] = next_logits[0, eos_id] / 4.0
+
+            next_token_id = self.sample_next_token(
+                logits=next_logits,
+                gen_tokens=gen_tokens,
+                seed_rng=sample_rng,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                repetition_penalty=repetition_penalty,
+                frequency_penalty=frequency_penalty,
+                no_repeat_ngram_size=no_repeat_ngram_size,
+                recent_tokens_window=200,
+            )
+
+            if next_token_id == eos_id:
+                break
+
+            next_tok_tensor = torch.tensor([[next_token_id]], dtype=torch.long).to(self.device)
+            gen_tokens = torch.cat([gen_tokens, next_tok_tensor], dim=1)
+            new_tokens_generated += 1
+            
+            if stream:
+                yield self.token_encoder.decode([next_token_id], errors='ignore')
+        
+        if not stream:
+            yield self.token_encoder.decode(gen_tokens[0, :].tolist(), errors='ignore')
+    
+    def print_stream(
+        self,
+        prompt,
+        max_new_tokens=200,
+        seed=42,
+        longer_story=True,
+        temperature=0.8,
+        top_k=None,
+        top_p=0.9,
+        repetition_penalty=1.2,
+        frequency_penalty=0.6,
+        no_repeat_ngram_size=3,
+        context_window=512,
+    ):
+        text = prompt
+        last_piece = ""
+        print(prompt, end="", flush=True)
+        for piece in self.generate(
+            prompt,
+            max_new_tokens=max_new_tokens,
+            seed=seed,
+            longer_story=longer_story,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+            frequency_penalty=frequency_penalty,
+            no_repeat_ngram_size=no_repeat_ngram_size,
+            context_window=context_window,
+        ):
+            if piece == last_piece:
+                continue
+            last_piece = piece
+            text = concat(text, piece)
+            print(piece, end="", flush=True)
+        return text

src/model.py

@@ -0,0 +1,179 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as f
+from dataclasses import dataclass
+import inspect
+
+@dataclass
+class Config:
+    context_length : int = 1024
+    vocab_size: int = 50257
+    num_layers : int = 12
+    embedding_dim : int = 768
+    num_heads: int = 12
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self,config : Config,masked=False):
+        super(MultiHeadAttention,self).__init__()
+        self.num_heads = config.num_heads
+        self.masked = masked
+        self.embedding_dim = config.embedding_dim
+        self.c_attention = nn.Linear(config.embedding_dim,3*config.embedding_dim)
+        self.c_projection = nn.Linear(config.embedding_dim,config.embedding_dim)
+        self.c_projection.SCALE_INIT = 1.0
+    
+    def forward(self,x):
+        B, T, C = x.shape
+        QKV = self.c_attention(x)
+        Query_q,Key_k,Value_v = QKV.split(self.embedding_dim,dim=-1)
+        Query_q = Query_q.view(B,T,self.num_heads,self.embedding_dim//self.num_heads).transpose(1,2)
+        Key_k = Key_k.view(B,T,self.num_heads,self.embedding_dim//self.num_heads).transpose(1,2)
+        Value_v = Value_v.view(B,T,self.num_heads,self.embedding_dim//self.num_heads).transpose(1,2)
+
+        # out = f.scaled_dot_product_attention(Query_q,Key_k,Value_v,is_causal=True)
+        if self.masked:
+            out = f.scaled_dot_product_attention(Query_q,Key_k,Value_v,is_causal=True)
+        else:
+            out = f.scaled_dot_product_attention(Query_q,Key_k,Value_v,is_causal=False)
+        out = out.transpose(1,2).contiguous().view(B,T,C)
+        return self.c_projection(out)
+    
+class MLP(nn.Module):
+    def __init__(self,config : Config):
+        super(MLP,self).__init__()
+        self.c_fc = nn.Linear(config.embedding_dim,4*config.embedding_dim)
+        self.gelu = nn.GELU(approximate='tanh')
+        self.c_projection = nn.Linear(4*config.embedding_dim,config.embedding_dim)
+        self.c_projection.SCALE_INIT = 1.0
+    def forward(self,x):
+        x = self.c_fc(x)
+        x = self.gelu(x)
+        x = self.c_projection(x)
+        return x
+    
+class DecoderBlock(nn.Module):
+    def __init__(self,config : Config):
+        """Decoder block without the encoder output"""
+        super(DecoderBlock,self).__init__()
+        self.masked_attention = MultiHeadAttention(config,masked=True)
+        self.layer_norm1 = nn.LayerNorm(config.embedding_dim)
+        # self.attention = MultiHeadAttention(config,masked=False)
+        # self.layer_norm2 = nn.LayerNorm(config.embedding_dim)
+        self.mlp = MLP(config)
+        self.layer_norm3 = nn.LayerNorm(config.embedding_dim)
+    
+    def forward(self,x):
+        x = x + self.masked_attention(self.layer_norm1(x))
+        # x = x + self.attention(self.layer_norm2(x))
+        x = x + self.mlp(self.layer_norm3(x))
+        return x
+
+class TransformerDecoder(nn.Module):
+    def __init__(self,config : Config):
+        super(TransformerDecoder,self).__init__()
+        self.config = config
+        self.word_token_embedding = nn.Embedding(self.config.vocab_size,self.config.embedding_dim)
+        self.word_position_embedding = nn.Embedding(self.config.context_length,self.config.embedding_dim)
+        layers = [DecoderBlock(config) for _ in range(config.num_layers)]
+        self.hidden_layers = nn.Sequential(*layers)
+        self.layer_norm = nn.LayerNorm(self.config.embedding_dim)
+    
+    def forward(self,idx):
+        B,T = idx.shape
+        pos = torch.arange(0,T,dtype=torch.long,device=idx.device)
+        pos_embed = self.word_position_embedding(pos)
+        token_embed = self.word_token_embedding(idx)
+        x = pos_embed + token_embed
+        x = self.hidden_layers(x)
+        x = self.layer_norm(x)
+        return x
+
+class GPT(nn.Module):
+    def __init__(self,config : Config):
+        super(GPT,self).__init__()
+        self.config=config
+        self.transformerDecoder = TransformerDecoder(config)
+        self.language_modeling_head = nn.Linear(config.embedding_dim,config.vocab_size,bias=False)
+        self.transformerDecoder.word_token_embedding.weight = self.language_modeling_head.weight
+        self.apply(self._init_weights)
+    
+    def _init_weights(self,module):
+        if isinstance(module,nn.Linear):
+            std=0.02
+            if hasattr(module,'SCALE_INIT'):
+                std /= (2*self.config.num_layers)**0.5
+            torch.nn.init.normal_(module.weight,mean=0,std=std)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module,nn.Embedding):
+            torch.nn.init.normal_(module.weight,mean=0,std=0.02)
+    
+    def forward(self,idx,targets=None):
+        x = self.transformerDecoder(idx)
+        logits = self.language_modeling_head(x)
+        loss = None
+        if targets is not None:
+            loss = f.cross_entropy(logits.view(-1,logits.shape[-1]),targets.view(-1))
+        return logits,loss  
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens=50, temperature=0.8, top_k=None, do_sample=False, eos_token_id=None):
+        self.eval()
+
+        B, T = idx.shape
+        device = idx.device
+        context_len = self.config.context_length
+
+        if T > context_len:
+            idx = idx[:, -context_len:]
+            T = idx.shape[1]
+
+        generated = idx.clone()
+
+        for _ in range(max_new_tokens):
+            input_ids = generated[:, -context_len:]
+
+            logits, _ = self.forward(input_ids, targets=None) 
+            next_logits = logits[:, -1, :] 
+
+            if temperature != 1.0 and temperature > 0.0:
+                next_logits = next_logits / temperature
+
+            if do_sample:
+                if top_k is not None and top_k > 0:
+                    vals, idxs = next_logits.topk(top_k, dim=-1)
+                    min_vals = vals[:, -1].unsqueeze(-1) 
+                    mask = next_logits < min_vals
+                    next_logits = next_logits.masked_fill(mask, float('-inf'))
+
+                probs = torch.softmax(next_logits, dim=-1) 
+                next_token = torch.multinomial(probs, num_samples=1)  
+            else:
+                next_token = torch.argmax(next_logits, dim=-1, keepdim=True) 
+
+            generated = torch.cat([generated, next_token], dim=1)  
+
+            if eos_token_id is not None:
+                if (generated == eos_token_id).any(dim=1).all():
+                    break
+
+        return generated
+    def configure_optimizer(self,weight_decay,lr,device_type,master_process):
+        param_dict = {pn:p for pn, p in self.named_parameters() if p.requires_grad}
+
+        decay_params = [p for pn, p in param_dict.items() if p.dim() >=2]
+        nodecay_params = [p for pn, p in param_dict.items() if p.dim() < 2]
+        optim_groups = [
+            {'params':decay_params,'weight_decay':weight_decay},
+            {'params':nodecay_params,'weight_decay':0.0}
+        ]
+        num_decay_params = sum(p.numel() for p in decay_params)
+        num_nodecay_params = sum(p.numel() for p in nodecay_params)
+        if master_process:
+            print(f'num decay parameter tensors: {len(decay_params)} with {num_decay_params:,} parameters')
+            print(f'num nodecay parameter tensors: {len(nodecay_params)} with {num_nodecay_params:,} parameters')
+        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
+        use_fused = fused_available and device_type == 'cuda'
+        if master_process:
+            print(f'using fused AdamW optimizer: {use_fused}')
+        optimizer = torch.optim.AdamW(optim_groups, lr=lr, betas=(0.9, 0.95), eps=1e-8, fused=use_fused)
+        return optimizer

src/trainer.py

@@ -0,0 +1,235 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as f
+from dataclasses import dataclass
+import time
+import os
+from src.model import GPT,Config
+from sacrebleu import corpus_bleu
+from rouge_score import rouge_scorer
+import numpy as np
+import math
+
+torch.set_float32_matmul_precision('high')
+
+def repetition_rate(text, n=3):
+    tokens = text.split()
+    if len(tokens) < n:
+        return 0.0
+    ngrams = [tuple(tokens[i:i+n]) for i in range(len(tokens)-n+1)]
+    return (len(ngrams) - len(set(ngrams))) / len(ngrams)
+
+def distinct_n(text, n=1):
+    tokens = text.split()
+    if len(tokens) < n:
+        return 0.0
+    
+    ngrams = [tuple(tokens[i:i+n]) for i in range(len(tokens) - n + 1)]
+    return len(set(ngrams)) / len(ngrams)
+
+
+def compute_self_bleu(generated_texts):
+    if len(generated_texts) < 2:
+        return 0.0
+
+    scores = []
+    N = len(generated_texts)
+
+    for i in range(N):
+        hyp = generated_texts[i]
+        refs = generated_texts[:i] + generated_texts[i+1:]
+
+        bleu = corpus_bleu([hyp], [refs]).score
+        scores.append(bleu)
+
+    return sum(scores) / len(scores)
+
+class Trainer:
+    def __init__(self,model : GPT,optimizer,train_loader,val_loader,token_encoder,eval_freq,grad_accum_steps,device,master_process,logpath):
+        self.model = model
+        self.optimizer = optimizer
+        self.train_loader = train_loader
+        self.val_loader = val_loader
+        self.token_encoder = token_encoder
+        self.master_process = master_process
+        self.eval_freq = eval_freq
+        self.grad_accum_steps = grad_accum_steps
+        self.device = device
+        self.device_type = 'cuda' if device.startswith('cuda') else 'cpu'
+        self.logpath=logpath
+
+    
+    def train(self,max_steps,warmup_steps,max_lr,min_lr):
+        history={
+            'val_losses':[],
+            'perplexities':[],
+            'train_losses':[]
+        }
+        for step in range(max_steps):
+            val_loss = None
+            perplexity=None
+            t0 = time.time()
+            self.is_last_step = (step == max_steps-1)
+            self.model.train()
+            self.optimizer.zero_grad()
+            batch_loss = 0.0
+
+            for mini_step in range(self.grad_accum_steps):
+                inp, target = self.train_loader.next_batch()
+                inp, target = inp.to(self.device),target.to(self.device)
+
+                with torch.autocast(device_type=self.device_type,dtype=torch.bfloat16):
+                    logits,loss = self.model(inp,target)
+                loss /=self.grad_accum_steps
+                batch_loss+=loss.detach()
+                loss.backward()
+            norm = nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
+            lr = self.estimate_lr(step,warmup_steps,max_steps,max_lr,min_lr)
+            for param_group in self.optimizer.param_groups:
+                param_group['lr'] = lr
+            self.optimizer.step()
+            if self.device_type == 'cuda':
+                torch.cuda.synchronize()
+            dt = (time.time() - t0) * 1000.0    # in ms
+            tokens_processed = self.train_loader.B * self.train_loader.T * self.grad_accum_steps * 1
+            tokens_per_sec = tokens_processed / dt
+            
+            if step % self.eval_freq == 0 or self.is_last_step:
+                val_loss,perplexity = self.evaluate_validation(step)
+                history['val_losses'].append(val_loss)
+                history['perplexities'].append(perplexity)
+            
+            history['train_losses'].append(batch_loss.item())
+            if self.master_process:
+                print(f'step {step:4d} | train loss: {batch_loss.item():.2f}{f' | val loss: {val_loss:.2f}' if val_loss is not None else ''}{f' | perplexity: {perplexity:.2f}' if perplexity is not None else ''} | lr: {lr:.2e} | norm: {norm:.4f} | dt: {dt:.4f}ms | tok/sec: {tokens_per_sec:.4f}')
+                with open(self.logpath, 'a') as f:
+                    f.write(f'{step} train {batch_loss.item():.6f}\n')
+
+        evaluation =self.evaluate_text_metrics(
+            max_samples=60,
+            gen_len=256,
+            do_sample=False,    
+            top_k=None,
+            temperature=0.2,
+            eos_token_id=None    
+        )
+        return history,evaluation
+    
+    def evaluate_text_metrics(self, max_samples=100, gen_len=50, do_sample=False, top_k=None, temperature=1.0, eos_token_id=None):
+        self.model.eval()
+        self.val_loader.reset()
+
+        hyps = []
+        refs = []
+        scorer = rouge_scorer.RougeScorer(['rougeL'], use_stemmer=True)
+
+        samples_collected = 0
+        while samples_collected < max_samples:
+            try:
+                inp, target = self.val_loader.next_batch()
+            except StopIteration:
+                break
+
+            inp = inp.to(self.device)
+            target = target.to(self.device)
+
+            if inp.shape[1] > self.model.config.context_length:
+                inp = inp[:, -self.model.config.context_length:]
+
+            with torch.no_grad():
+                generated = self.model.generate(
+                    inp,
+                    max_new_tokens=gen_len,
+                    temperature=temperature,
+                    top_k=top_k,
+                    do_sample=do_sample,
+                    eos_token_id=eos_token_id
+                )
+
+            B = generated.shape[0]
+            for i in range(B):
+                gen_ids = generated[i, inp.shape[1]:].tolist()
+
+                pred_text = self.token_encoder.decode(gen_ids)
+                ref_text  = self.token_encoder.decode(target[i].tolist())
+
+                hyps.append(pred_text)
+                refs.append(ref_text)
+                samples_collected += 1
+                if samples_collected >= max_samples:
+                    break
+
+        if len(hyps) == 0:
+            return 0.0, 0.0
+
+        rep_scores = []
+        distinct1_scores = []
+        distinct2_scores = []
+
+        for txt in hyps:
+            rep_scores.append(repetition_rate(txt, n=3))
+            distinct1_scores.append(distinct_n(txt, n=1))
+            distinct2_scores.append(distinct_n(txt, n=2))
+
+        avg_rep = sum(rep_scores) / len(rep_scores)
+        avg_d1 = sum(distinct1_scores) / len(distinct1_scores)
+        avg_d2 = sum(distinct2_scores) / len(distinct2_scores)
+
+
+        bleu = corpus_bleu(hyps, refs).score
+        self_bleu=compute_self_bleu(hyps)
+        rouge_scores = []
+        for h, r in zip(hyps, refs):
+            sc = scorer.score(r, h)['rougeL'].fmeasure
+            rouge_scores.append(sc)
+        rouge_l = sum(rouge_scores) / len(rouge_scores)
+
+        if self.master_process:
+            print(f"[Text Eval] samples={len(hyps)} BLEU={bleu:.2f} ROUGE-L={rouge_l:.4f} SELF-BLEU={self_bleu:.2f} REP={avg_rep:.4f} D1={avg_d1:.4f} D2={avg_d2:.4f}")
+            with open(self.logpath, 'a') as f:
+                f.write(f"eval samples={len(hyps)} BLEU={bleu:.2f} ROUGE-L={rouge_l:.4f} SELF-BLEU={self_bleu:.2f} REP={avg_rep:.4f} D1={avg_d1:.4f} D2={avg_d2:.4f}\n")
+
+        return {"bleu":bleu,"rogue-l":rouge_scores,"self-bleu":self_bleu,"repetition":rep_scores,"D1":distinct1_scores,"D2":distinct2_scores}
+
+    def evaluate_validation(self,step):
+        self.model.eval()
+        self.val_loader.reset()
+        with torch.no_grad():
+            val_loss_accum = 0.0
+            val_steps = 20
+            for _ in range(val_steps):
+                inp, target = self.val_loader.next_batch()
+                inp, target = inp.to(self.device),target.to(self.device)
+
+                with torch.autocast(device_type=self.device_type,dtype=torch.bfloat16):
+                    logits,loss = self.model(inp,target)
+                loss /=val_steps
+                val_loss_accum+=loss.detach()
+
+        if self.master_process:
+            perplexity = math.exp(val_loss_accum.item())
+            with open(self.logpath, 'a') as f:
+                f.write(f'{step} val {val_loss_accum.item():.4f}\n')
+
+            if step > 0 and (step % 10000 == 0 or self.is_last_step):
+                raw_model = self.model
+                logdir = os.path.dirname(self.logpath)
+                ckpt_path = os.path.join(logdir, f'model_{step:05d}.pt')
+                checkpoint = {
+                    'model': raw_model.state_dict(),
+                    'config': raw_model.config,
+                    'step': step,
+                    'val_loss': val_loss_accum.item()
+                }
+                torch.save(checkpoint, ckpt_path)
+        return val_loss_accum.item(),perplexity
+    
+    def estimate_lr(self, step, warmup_steps, max_steps, max_lr, min_lr):
+        if step < warmup_steps:
+            return max_lr * (step+1) / warmup_steps
+        if step > max_steps:
+            return min_lr
+        decay_ratio = (step - warmup_steps) / (max_steps - warmup_steps)
+        assert 0 <= decay_ratio <= 1
+        coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
+        return min_lr + coeff * (max_lr - min_lr)