Upload app.py

app.py

@@ -0,0 +1,99 @@
+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 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()