generation_example.py

#!/usr/bin/env python3
"""Example script for text generation with a trained model."""

import argparse
import sys
from pathlib import Path

# Add src to path
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

import torch
from lm_trainer.inference.generation import generate_text
from lm_trainer.models.architectures import TransformerModel
from lm_trainer.training.configuration import ModelConfig
from lm_trainer.tokenizers import BPETokenizer


def main():
    """Generate text with a trained model."""
    
    parser = argparse.ArgumentParser(description="Generate text with trained model")
    parser.add_argument("--prompt", type=str, default="The future of AI is",
                       help="Input prompt for text generation")
    parser.add_argument("--max-length", type=int, default=100,
                       help="Maximum length of generated text")
    parser.add_argument("--temperature", type=float, default=0.8,
                       help="Sampling temperature")
    parser.add_argument("--top-k", type=int, default=50,
                       help="Top-k sampling parameter")
    parser.add_argument("--top-p", type=float, default=0.95,
                       help="Top-p (nucleus) sampling parameter")
    parser.add_argument("--generation-method", type=str, default="sampling",
                       choices=["greedy", "beam_search", "sampling"],
                       help="Text generation method")
    
    args = parser.parse_args()
    
    print("Loading model and tokenizer...")
    
    # Create model configuration
    model_config = ModelConfig(
        vocab_size=32000,
        d_model=256,
        n_heads=4,
        n_layers=4,
        d_ff=1024,
        max_seq_len=512
    )
    
    # Create model
    model = TransformerModel(model_config)
    
    # Create tokenizer
    tokenizer = BPETokenizer()
    # For demo, create a simple tokenizer
    tokenizer.train_from_texts([
        "The quick brown fox jumps over the lazy dog",
        "Machine learning is a subset of artificial intelligence",
        "Natural language processing enables human-computer interaction",
        "The future of AI is", "exciting", "promising", "challenging"
    ])
    
    # In a real scenario, you would load a trained model:
    # model.load_state_dict(torch.load("path/to/model.pt"))
    
    print(f"Generating text with method: {args.generation_method}")
    print(f"Prompt: {args.prompt}")
    
    # Generate text
    generated_text = generate_text(
        model=model,
        tokenizer=tokenizer,
        prompt=args.prompt,
        generation_method=args.generation_method,
        max_length=args.max_length,
        temperature=args.temperature,
        top_k=args.top_k,
        top_p=args.top_p
    )
    
    print("\n" + "="*50)
    print("GENERATED TEXT")
    print("="*50)
    print(generated_text)
    print("="*50)


if __name__ == "__main__":
    main()