import os
import sys
import torch
import torch.nn.functional as F

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from model.model import ModelConfig, Transformer

def generate(model: Transformer, prompt_tokens: torch.Tensor, max_new_tokens: int = 50, temperature: float = 0.8, top_k: int = 40):
    model.eval()
    tokens = prompt_tokens.clone()
    
    with torch.no_grad():
        for _ in range(max_new_tokens):
            logits = model(tokens)
            logits = logits[:, -1, :] / temperature
            
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('Inf')
            
            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            tokens = torch.cat((tokens, next_token), dim=1)
            
    return tokens

if __name__ == "__main__":
    device = "cuda" if torch.cuda.is_available() else "cpu"
    config = ModelConfig.get_125m()
    model = Transformer(config).to(device)
    
    dummy_input = torch.randint(0, config.vocab_size, (1, 10), device=device)
    print("Generating sample sequence from model...")
    output = generate(model, dummy_input, max_new_tokens=20)
    print(f"Generated Tokens Shape: {output.shape}")
    print(f"Generated Token IDs: {output[0].tolist()}")