Spaces:
Running
Running
| 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-025000-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 fn_query_on_load(): | |
| return "Biofuels would disrupt" | |
| def generate_output(prompt, max_new_tokens=200, temperature=0.8, top_k=50): | |
| m.eval() | |
| encoded_text = tokenizer_gpt.encode(prompt) | |
| # print('--------------------encoded text = ',encoded_text) | |
| 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=max_new_tokens, temperature=0.8, top_k=50)[0]) | |
| m.train() | |
| return { | |
| output: out_text | |
| } | |
| with gr.Blocks() as app: | |
| with gr.Row(): | |
| gr.Markdown( | |
| """ | |
| # MiniGPT - GPT Training on LLaMa with redpajama dataset | |
| ### Enter a context to generate automated text " | |
| """) | |
| with gr.Row(visible=True): | |
| search_text = gr.Textbox(value=fn_query_on_load, placeholder='Enter prompt..', label='Enter Prompt') | |
| with gr.Row(): | |
| submit_btn = gr.Button("Submit", variant='primary') | |
| clear_btn = gr.ClearButton() | |
| with gr.Row(): | |
| with gr.Row(): | |
| output = gr.Textbox(lines=15, interactive=False, label='Out Box') | |
| def clear_data(): | |
| return { | |
| output: None, | |
| search_text: None | |
| } | |
| clear_btn.click(clear_data, None, [output, search_text]) | |
| submit_btn.click( | |
| generate_output, | |
| search_text, | |
| output | |
| ) | |
| ''' | |
| Launch the app | |
| ''' | |
| app.queue().launch() |