Spaces:
Running
Running
import gradio as gr | |
from transformers import pipeline | |
# Ensure the correct Keras package is used | |
import tensorflow as tf | |
import tf_keras | |
# Define the model and the text generation function | |
fine_tuned_model = "Amitesh007/text_generation-finetuned-gpt2" | |
generator = pipeline('text-generation', model=fine_tuned_model) | |
def generate(text): | |
results = generator(text, num_return_sequences=2, max_length=100) | |
return results[0]["generated_text"], results[1]["generated_text"] | |
# Create the Gradio Blocks interface | |
with gr.Blocks() as demo: | |
gr.Markdown("# Text Generator GPT2 Pipeline") | |
gr.Markdown("This is a fine-tuned base GPT2 model inference, trained on a small 'Game of Thrones' dataset.") | |
with gr.Row(): | |
with gr.Column(): | |
input_text = gr.Textbox(lines=5, label="Input Text here....", placeholder="Type a sentence to start generating text") | |
generate_button = gr.Button("Generate") | |
with gr.Column(): | |
output_text1 = gr.Textbox(label="Generated Text 1") | |
output_text2 = gr.Textbox(label="Generated Text 2") | |
examples = gr.Examples( | |
examples=[["A light snow had fallen the night before, and there were"], | |
["The pig face had been smashed in with a mace, but Tyrion"]], | |
inputs=input_text | |
) | |
generate_button.click(fn=generate, inputs=input_text, outputs=[output_text1, output_text2], concurrency_limit=2) | |
# Launch the interface with the max_threads parameter to control the total number of workers | |
demo.launch(max_threads=8) | |