Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| import gradio as gr | |
| from tsai_gpt.tokenizer import Tokenizer | |
| import lightning as L | |
| from lightning.fabric.loggers import CSVLogger | |
| from pathlib import Path | |
| from tsai_gpt.utils import num_parameters, load_checkpoint, get_default_supported_precision | |
| from tsai_gpt.model import GPT, Block, Config | |
| model_name = "pythia-160m" | |
| name = "redpajama" | |
| out_dir = Path("out") / name | |
| log_interval = 100 | |
| precision = get_default_supported_precision(False) | |
| logger = CSVLogger("out", name, flush_logs_every_n_steps=log_interval) | |
| fabric = L.Fabric(devices=1, strategy="auto", precision=precision, loggers=logger) | |
| config = Config.from_name(model_name) | |
| def _init_weights(module: nn.Module) -> None: | |
| """Meant to be used with `gpt.apply(gpt._init_weights)`.""" | |
| if isinstance(module, nn.Linear): | |
| torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) | |
| 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.0, std=0.02) | |
| with fabric.init_module(empty_init=True): | |
| model = GPT(config) | |
| model.apply(_init_weights) | |
| model.apply(_init_weights) | |
| checkpoint_path = Path("out/redpajama/iter-015000-ckpt.pth") | |
| load_checkpoint(fabric, model, checkpoint_path) | |
| #print(model.transformer.h[0].mlp.fc.weight) | |
| #fabric.print(f"Time to instantiate model: {time.perf_counter() - t0:.02f} seconds.") | |
| #fabric.print(f"Total parameters {num_parameters(model):,}") | |
| weight_decay = 1e-1 | |
| beta1 = 0.9 | |
| beta2 = 0.95 | |
| learning_rate = 6e-3 | |
| hparams = {k: v for k, v in locals().items() if isinstance(v, (int, float, str)) and not k.startswith("_")} | |
| model = fabric.setup(model) | |
| optimizer = torch.optim.AdamW( | |
| model.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=(beta1, beta2), foreach=False | |
| ) | |
| # model_copy = model | |
| optimizer = fabric.setup_optimizers(optimizer) | |
| state = {"model": model, "optimizer": optimizer, "hparams": hparams, "iter_num": 0, "step_count": 0} | |
| resume = max(out_dir.glob("*.pth"), key=lambda p: int(p.name.split("-")[1])) | |
| if resume: | |
| fabric.print(f"Loading model from {resume}") | |
| fabric.load(resume, state) | |
| deviceType = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| m = model.to(deviceType) | |
| tokenizer_gpt = Tokenizer(checkpoint_dir=Path("checkpoints/meta-llama/Llama-2-7b-chat-hf")) | |
| def inference(input_context, count): | |
| #print('--------------------input = ',input_context) | |
| encoded_text = tokenizer_gpt.encode(input_context) | |
| #print('--------------------encoded text = ',encoded_text) | |
| count = int(count) | |
| #print('--------------------count = ',count) | |
| reshaped_tensor = torch.unsqueeze(encoded_text, 0).to(deviceType) | |
| #print('--------------------reshaped_tensor = ',reshaped_tensor) | |
| out_text = tokenizer_gpt.decode(m.generate(reshaped_tensor, max_new_tokens=count)[0]) | |
| return out_text | |
| title = "TSAI S22 Assignment: GPT training on LLaMa - redpajama dataset" | |
| description = "A simple Gradio interface that accepts a context and generates text " | |
| examples = [["Machine Learning","200"], | |
| ["Deep Learning","200"] | |
| ] | |
| demo = gr.Interface( | |
| inference, | |
| inputs = [gr.Textbox(placeholder="Enter starting characters"), gr.Textbox(placeholder="Enter number of characters you want to generate")], | |
| outputs = [gr.Textbox(label="Generated text")], | |
| title = title, | |
| description = description, | |
| examples = examples | |
| ) | |
| demo.launch() |