ToletiSri commited on
Commit
91a9739
1 Parent(s): 36bf165

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +99 -0
app.py ADDED
@@ -0,0 +1,99 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ import gradio as gr
4
+ from tsai_gpt.tokenizer import Tokenizer
5
+ import lightning as L
6
+ from lightning.fabric.loggers import CSVLogger
7
+ from pathlib import Path
8
+ from tsai_gpt.utils import num_parameters, load_checkpoint, get_default_supported_precision
9
+ from tsai_gpt.model import GPT, Block, Config
10
+
11
+ model_name = "pythia-160m"
12
+ name = "redpajama"
13
+ out_dir = Path("out") / name
14
+ log_interval = 100
15
+
16
+ precision = get_default_supported_precision(False)
17
+ logger = CSVLogger("out", name, flush_logs_every_n_steps=log_interval)
18
+ fabric = L.Fabric(devices=1, strategy="auto", precision=precision, loggers=logger)
19
+
20
+ config = Config.from_name(model_name)
21
+
22
+ def _init_weights(module: nn.Module) -> None:
23
+ """Meant to be used with `gpt.apply(gpt._init_weights)`."""
24
+ if isinstance(module, nn.Linear):
25
+ torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
26
+ if module.bias is not None:
27
+ torch.nn.init.zeros_(module.bias)
28
+ elif isinstance(module, nn.Embedding):
29
+ torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
30
+
31
+ with fabric.init_module(empty_init=True):
32
+ model = GPT(config)
33
+ model.apply(_init_weights)
34
+ model.apply(_init_weights)
35
+
36
+
37
+ checkpoint_path = Path("out/redpajama/iter-015000-ckpt.pth")
38
+
39
+ load_checkpoint(fabric, model, checkpoint_path)
40
+
41
+ #print(model.transformer.h[0].mlp.fc.weight)
42
+
43
+ #fabric.print(f"Time to instantiate model: {time.perf_counter() - t0:.02f} seconds.")
44
+ #fabric.print(f"Total parameters {num_parameters(model):,}")
45
+
46
+ weight_decay = 1e-1
47
+ beta1 = 0.9
48
+ beta2 = 0.95
49
+ learning_rate = 6e-3
50
+ hparams = {k: v for k, v in locals().items() if isinstance(v, (int, float, str)) and not k.startswith("_")}
51
+
52
+ model = fabric.setup(model)
53
+ optimizer = torch.optim.AdamW(
54
+ model.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=(beta1, beta2), foreach=False
55
+ )
56
+
57
+ # model_copy = model
58
+
59
+ optimizer = fabric.setup_optimizers(optimizer)
60
+
61
+ state = {"model": model, "optimizer": optimizer, "hparams": hparams, "iter_num": 0, "step_count": 0}
62
+
63
+ resume = max(out_dir.glob("*.pth"), key=lambda p: int(p.name.split("-")[1]))
64
+ if resume:
65
+ fabric.print(f"Loading model from {resume}")
66
+ fabric.load(resume, state)
67
+
68
+ deviceType = 'cuda' if torch.cuda.is_available() else 'cpu'
69
+ m = model.to(deviceType)
70
+ tokenizer_gpt = Tokenizer(checkpoint_dir=Path("checkpoints\meta-llama\Llama-2-7b-chat-hf"))
71
+
72
+
73
+ def inference(input_context, count):
74
+ #print('--------------------input = ',input_context)
75
+ encoded_text = tokenizer_gpt.encode(input_context)
76
+ #print('--------------------encoded text = ',encoded_text)
77
+ count = int(count)
78
+ #print('--------------------count = ',count)
79
+ reshaped_tensor = torch.unsqueeze(encoded_text, 0).to(deviceType)
80
+ #print('--------------------reshaped_tensor = ',reshaped_tensor)
81
+ out_text = tokenizer_gpt.decode(m.generate(reshaped_tensor, max_new_tokens=count)[0])
82
+ return out_text
83
+
84
+ title = "TSAI S22 Assignment: GPT training on LLaMa dataset"
85
+ description = "A simple Gradio interface that accepts a context and generates text "
86
+ examples = [["Machine Learning","200"],
87
+ ["Deep Learning","200"]
88
+ ]
89
+
90
+
91
+ demo = gr.Interface(
92
+ inference,
93
+ inputs = [gr.Textbox(placeholder="Enter starting characters"), gr.Textbox(placeholder="Enter number of characters you want to generate")],
94
+ outputs = [gr.Textbox(label="Generated text")],
95
+ title = title,
96
+ description = description,
97
+ examples = examples
98
+ )
99
+ demo.launch()