|
from transformers import BartForConditionalGeneration, BartTokenizer |
|
from transformers.utils import logging |
|
import gradio as gr |
|
|
|
|
|
logger = logging.get_logger("transformers") |
|
|
|
|
|
model_name = "facebook/bart-large-cnn" |
|
tokenizer = BartTokenizer.from_pretrained(model_name) |
|
model = BartForConditionalGeneration.from_pretrained(model_name) |
|
|
|
|
|
def generate_text(prompt): |
|
inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512) |
|
summary_ids = model.generate(inputs["input_ids"], max_length=150, min_length=40, length_penalty=2.0, num_beams=4, early_stopping=True) |
|
return tokenizer.decode(summary_ids[0], skip_special_tokens=True) |
|
|
|
|
|
|
|
|
|
|
|
|
|
def clear_save_textbox(message): |
|
return " ", message |
|
|
|
def show_input_text(message,history:list[tuple[str,str]]): |
|
history.append((message,"")) |
|
story = generate_text(message) |
|
history[-1] = (message,story) |
|
return history |
|
|
|
def delete_previous_text(history:list[tuple[str,str]]): |
|
try: |
|
message, _ = history.pop() |
|
except IndexError: |
|
message = " " |
|
return history, message |
|
|
|
|
|
interface = gr.Interface(fn=generate_text, inputs="text", outputs="text",title="TeLLMyStory",description="Enter your story idea and the model will generate the story based on it.") |
|
with gr.Blocks() as demo: |
|
gr.Markdown("TeLLMyStory chatbot") |
|
|
|
with gr.Row(): |
|
input_text = gr.Textbox(label="Enter your story idea here") |
|
|
|
|
|
|
|
|
|
with gr.Row(): |
|
gr.Markdown("History of your story ideas") |
|
gen_story = gr.Textbox(label="History") |
|
|
|
|
|
|
|
|
|
interface.launch() |
|
|