| import gradio as gr |
| import torch |
| import tiktoken |
| from src.model import TinyLLM, Config |
|
|
| |
| config = Config() |
| config.block_size = 128 |
| enc = tiktoken.get_encoding("gpt2") |
|
|
| |
| model = TinyLLM(config) |
|
|
| |
| import os |
| ckpt_paths = ["checkpoints/tinyllm.pt", "tinyllm.pt", "checkpoints/best.pt"] |
| ckpt = None |
| for p in ckpt_paths: |
| if os.path.exists(p): |
| ckpt = p |
| break |
|
|
| if ckpt is None: |
| print(f"No checkpoint found, tried: {ckpt_paths}") |
| print("Training one first or downloading from HF...") |
| from huggingface_hub import hf_hub_download |
| ckpt = hf_hub_download("grahmseinstein/tinyllm-28m", "tinyllm.pt", local_dir=".") |
|
|
| print(f"Loading checkpoint: {ckpt}") |
| model.load_state_dict(torch.load(ckpt, map_location='cpu')) |
| model.eval() |
|
|
| print(f"Model loaded: {sum(p.numel() for p in model.parameters())/1e6:.2f}M params") |
|
|
| def generate(prompt, max_tokens): |
| if not prompt or prompt.strip() == "": |
| prompt = "Hello" |
| |
| ids = torch.tensor([enc.encode(prompt)], dtype=torch.long) |
| |
| with torch.no_grad(): |
| out = model.generate(ids, max_new_tokens=int(max_tokens)) |
| |
| text = enc.decode(out[0].tolist()) |
| return text |
|
|
| |
| with gr.Blocks(title="TinyLLM 28M - grahmseinstein") as demo: |
| gr.Markdown("# TinyLLM 28M Prototype\nby grahmseinstein - GPT2 50257 tokens") |
| with gr.Row(): |
| prompt = gr.Textbox(label="Prompt", value="Hello, I am a", lines=2) |
| max_tokens = gr.Slider(20, 200, value=100, step=10, label="Max tokens") |
| output = gr.Textbox(label="Generated", lines=10) |
| btn = gr.Button("Generate") |
| btn.click(fn=generate, inputs=[prompt, max_tokens], outputs=output) |
| prompt.submit(fn=generate, inputs=[prompt, max_tokens], outputs=output) |
|
|
| demo.launch(server_name="0.0.0.0", server_port=7860, share=False) |
|
|
