Spaces:
Sleeping
Sleeping
| import torch | |
| import tiktoken | |
| import gradio as gr | |
| import torch.nn.functional as F | |
| from model import GPT, GPTConfig | |
| device = 'cpu' | |
| if torch.cuda.is_available(): | |
| device = 'cuda' | |
| elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): | |
| device = "mps" | |
| model = GPT(GPTConfig()) | |
| ckpt = torch.load("gpt2.pt", map_location=torch.device(device)) | |
| unwanted_prefix = '_orig_mod.' | |
| for k,v in list(ckpt.items()): | |
| if k.startswith(unwanted_prefix): | |
| ckpt[k[len(unwanted_prefix):]] = ckpt.pop(k) | |
| model.load_state_dict(ckpt) | |
| model.to(device) | |
| enc = tiktoken.get_encoding('gpt2') | |
| def inference(input_text, num_return_sequences, max_length): | |
| input_tokens = torch.tensor(enc.encode(input_text), dtype=torch.long) | |
| input_tokens = input_tokens.unsqueeze(0).repeat(num_return_sequences, 1) | |
| x = input_tokens.to(device) | |
| while x.size(1) < max_length: | |
| # forward the model to get the logits | |
| with torch.no_grad(): | |
| logits = model(x)[0] # (B, T, vocab_size) | |
| # take the logits at the last position | |
| logits = logits[:, -1, :] # (B, vocab_size) | |
| # get the probabilities | |
| probs = F.softmax(logits, dim=-1) | |
| # do top-k sampling of 50 (huggingface pipeline default) | |
| # topk_probs here becomes (5, 50), topk_indices is (5, 50) | |
| topk_probs, topk_indices = torch.topk(probs, 50, dim=-1) | |
| # select a token from the top-k probabilities | |
| # note: multinomial does not demand the input to sum to 1 | |
| ix = torch.multinomial(topk_probs, 1) # (B, 1) | |
| # gather the corresponding indices | |
| xcol = torch.gather(topk_indices, -1, ix) # (B, 1) | |
| # append to the sequence | |
| x = torch.cat((x, xcol), dim=1) | |
| decode_list = [] | |
| # print the generated text | |
| for i in range(num_return_sequences): | |
| tokens = x[i, :max_length].tolist() | |
| decoded = enc.decode(tokens) | |
| decode_list.append(decoded) | |
| output = "\n======\n".join(decode_list) | |
| return output | |
| title = "GPT-2 trained on Shakespeare Plays dataset" | |
| description = "A simple Gradio interface to generate text from GPT-2 model trained on Shakespeare Plays" | |
| examples = [["Please put on these earmuffs because I can't hear you.", 5, 50], | |
| ["Twin 4-month-olds slept in the shade of the palm tree while the mother tanned in the sun.", 5, 50], | |
| ["Happiness can be found in the depths of chocolate pudding.", 5, 50], | |
| ["Seek success, but always be prepared for random cats.", 5, 50], | |
| ["This made him feel like an old-style rootbeer float smells.", 5, 50], | |
| ["The view from the lighthouse excited even the most seasoned traveler.", 5, 50], | |
| ["I've always wanted to go to Tajikistan, but my cat would miss me.", 5, 50], | |
| ["He found rain fascinating yet unpleasant.", 5, 50], | |
| ["Plans for this weekend include turning wine into water.", 5, 50], | |
| ["Iron pyrite is the most foolish of all minerals.", 5, 50], | |
| ] | |
| demo = gr.Interface( | |
| inference, | |
| inputs = [ | |
| gr.Textbox(label="Enter some text", type="text"), | |
| gr.Slider(minimum=1, maximum=5, step=1, value=5, label="Number of outputs"), | |
| gr.Slider(minimum=10, maximum=100, step=1, value=50, label="Maximum lenght of a sequence") | |
| ], | |
| outputs = [ | |
| gr.Textbox(label="Output", type="text") | |
| ], | |
| title = title, | |
| description = description, | |
| examples = examples, | |
| ) | |
| demo.launch() |