File size: 1,438 Bytes
38e3aab |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
tokenizer = AutoTokenizer.from_pretrained("fubuki119/JokesGPT")
model = AutoModelForCausalLM.from_pretrained("fubuki119/JokesGPT")
def generate(max_length):
starting_text = "JOKE:"
end_token = "<|endoftext|>"
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
model.eval()
with torch.no_grad():
cur_ids = torch.tensor(tokenizer.encode("JOKE:")
).unsqueeze(0).to(device)
for i in range(max_length):
outputs = model(cur_ids)
logits, _ = outputs[:]
softmax_logits = torch.softmax(logits[0, -1], dim=0)
next_token_id = torch.multinomial(softmax_logits, 1).item()
cur_ids = torch.cat([cur_ids, torch.ones(
(1, 1)).long().to(device) * next_token_id], dim=1)
if next_token_id == tokenizer.encode(end_token)[0]:
joke = cur_ids.detach().cpu().tolist()
joke = joke[0]
return tokenizer.decode(joke[3:-1])
joke = cur_ids.detach().cpu().tolist()
joke = joke[0]
return tokenizer.decode(joke[3:])
def fn(n):
print(f"i got {n}")
return "thanks"
iface = gr.Interface(
fn=generate,
inputs=gr.Number(value=200, label="Maxlen"),
outputs="text",
)
iface.launch(share=True, debug=True)
|