|
from transformers import AutoTokenizer, AutoModelForCausalLM |
|
import torch |
|
import gradio as gr |
|
|
|
tokenizer = AutoTokenizer.from_pretrained('TrLOX/gpt2-tdk') |
|
model = AutoModelForCausalLM.from_pretrained('TrLOX/gpt2-tdk') |
|
|
|
def text_generation(input_text, seed): |
|
input_ids = tokenizer(input_text, return_tensors="pt").input_ids |
|
torch.manual_seed(seed) |
|
outputs = model.generate(input_ids, do_sample=True, min_length=50, max_length=200) |
|
generated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True) |
|
return generated_text |
|
|
|
title = "TDK GPT2" |
|
description = "Title and description generation by keywords" |
|
|
|
gr.Interface( |
|
text_generation, |
|
[gr.inputs.Textbox(lines=2, label="Enter input text"), gr.inputs.Number(default=10, label="Enter seed number")], |
|
[gr.outputs.Textbox(type="auto", label="Text Generated")], |
|
title=title, |
|
description=description, |
|
theme="huggingface" |
|
).launch() |