Spaces:
Sleeping
Sleeping
| import torch | |
| import tiktoken | |
| from model import * | |
| import gradio as gr | |
| enc = tiktoken.get_encoding('gpt2') | |
| model = torch.load('model.pt',map_location='cpu') | |
| def response(message = "Hello, I'm a language model", num_return_sequences = 5,max_length = 30,top_k = 50): | |
| tokens = enc.encode(message) | |
| tokens = torch.tensor(tokens, dtype= torch.long) # (8,) #check tiktoken app | |
| tokens = tokens.unsqueeze(0).repeat(num_return_sequences, 1) # (5, 8) | |
| x = tokens.to('cpu') | |
| torch.manual_seed(42) | |
| torch.cuda.manual_seed(42) | |
| 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, top_k, 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) | |
| # print the generated text | |
| return_text = "" | |
| for i in range(num_return_sequences): | |
| tokens = x[i, :max_length].tolist() | |
| decoded = enc.decode(tokens) | |
| return_text = return_text + ">"+ decoded +"\n" | |
| return return_text | |
| # Function to generate text | |
| def generate_text(prompt, top_k, return_sequences, max_tokens): | |
| ans = response(message = prompt, num_return_sequences = return_sequences,max_length = max_tokens,top_k = top_k) | |
| return ans | |
| # Create Gradio interface | |
| iface = gr.Interface( | |
| fn=generate_text, | |
| inputs=[ | |
| gr.Textbox(lines=5, label="Input Text"), | |
| gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top-k"), | |
| gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Return Sequences"), | |
| gr.Slider(minimum=10, maximum=150, value=20, step=5, label="Max Tokens") | |
| ], | |
| outputs=gr.Textbox(label="Generated Text"), | |
| title="GPT-2 Text Generator", | |
| description="Generate text using GPT-2 model with adjustable parameters.", | |
| examples=[["Hello, I'm a language model"]], | |
| ) | |
| # Launch the interface | |
| iface.launch() | |