| """ |
| generate.py |
| =========== |
| Interactive text generation with the trained MiniLM model. |
| |
| Type a prompt and the model will complete it. |
| Type 'quit' or press Ctrl+C to exit. |
| |
| Author : André Costa |
| License : MIT |
| |
| Usage: |
| python3 generate.py |
| python3 generate.py --max-tokens 100 |
| python3 generate.py --temperature 0.9 --top-k 50 |
| """ |
|
|
| import argparse |
| import torch |
| from transformer import MiniLM, ModelConfig |
| from bpe_tokenizer import BPETokenizer |
|
|
|
|
| def load_model(checkpoint_path: str, tokenizer_path: str): |
| """Load the trained model and tokenizer.""" |
|
|
| print("Loading tokenizer...") |
| tokenizer = BPETokenizer.load(tokenizer_path) |
|
|
| print("Loading model...") |
| ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=True) |
|
|
| cfg_dict = ckpt["model_config"] |
| cfg_dict.pop("d_head", None) |
| config = ModelConfig(**cfg_dict) |
|
|
| model = MiniLM(config) |
|
|
| state_dict = ckpt["model_state"] |
| if any(k.startswith("_orig_mod.") for k in state_dict): |
| state_dict = {k.replace("_orig_mod.", ""): v for k, v in state_dict.items()} |
|
|
| model.load_state_dict(state_dict) |
| model.eval() |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| model = model.to(device) |
|
|
| print(f"Model ready — {config.n_params / 1e6:.1f}M parameters | device: {device}") |
| print(f"Vocab: {config.vocab_size} tokens | Context: {config.seq_len} tokens\n") |
|
|
| return model, tokenizer, device |
|
|
|
|
| def generate( |
| model, |
| tokenizer, |
| device, |
| prompt: str, |
| max_new_tokens: int, |
| temperature: float, |
| top_k: int, |
| top_p: float, |
| ) -> str: |
| """Generate text from a prompt.""" |
| input_ids = torch.tensor( |
| [tokenizer.encode(prompt)], |
| dtype=torch.long, |
| device=device, |
| ) |
|
|
| with torch.no_grad(): |
| output = model.generate( |
| input_ids, |
| max_new_tokens=max_new_tokens, |
| temperature=temperature, |
| top_k=top_k, |
| top_p=top_p, |
| ) |
|
|
| return tokenizer.decode(output[0].tolist()) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="MiniLM — Interactive text generation") |
| parser.add_argument("--checkpoint", type=str, default="./checkpoints/best_model.pt") |
| parser.add_argument("--tokenizer", type=str, default="./tokenizer") |
| parser.add_argument("--max-tokens", type=int, default=80) |
| parser.add_argument("--temperature", type=float, default=0.8) |
| parser.add_argument("--top-k", type=int, default=50) |
| parser.add_argument("--top-p", type=float, default=0.9) |
| args = parser.parse_args() |
|
|
| model, tokenizer, device = load_model(args.checkpoint, args.tokenizer) |
|
|
| print("=" * 55) |
| print(" MiniLM — Text Generation") |
| print(" Type a prompt and press Enter.") |
| print(" Type 'quit' to exit.") |
| print("=" * 55) |
| print() |
|
|
| while True: |
| try: |
| prompt = input("Prompt: ").strip() |
| except (KeyboardInterrupt, EOFError): |
| print("\nGoodbye!") |
| break |
|
|
| if not prompt: |
| continue |
|
|
| if prompt.lower() in ("quit", "exit", "q"): |
| print("Goodbye!") |
| break |
|
|
| result = generate( |
| model, tokenizer, device, |
| prompt=prompt, |
| max_new_tokens=args.max_tokens, |
| temperature=args.temperature, |
| top_k=args.top_k, |
| top_p=args.top_p, |
| ) |
|
|
| print(f"\n{result}\n") |
| print("-" * 55) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
