Spaces:
Sleeping
Sleeping
File size: 2,658 Bytes
b5f56f6 dcc2428 b5f56f6 c84d121 b5f56f6 c84d121 dcc2428 | 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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | # from transformers import pipeline
# import gradio as gr
# pipe = pipeline("translation", model="Helsinki-NLP/opus-mt-en-es")
# demo = gr.Interface.from_pipeline(pipe)
# demo.launch()
import gradio as gr
from transformers import pipeline
import torch
# Define the model to use.
MODEL_NAME = "google/flan-t5-small"
# Set up the pipeline, specifying the task and the model.
# The pipeline handles tokenization and model inference.
# Using device="cuda" if a GPU is available, otherwise falls back to CPU.
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = pipeline("text2text-generation", model=MODEL_NAME, device=device)
# Define the function that constructs the prompt and calls the pipeline.
def generate_text(user_input, prompt_template):
"""
Combines user input with a template and calls the Hugging Face transformers pipeline.
"""
# Create the full prompt based on the template and user input.
full_prompt = prompt_template.format(user_input=user_input)
# Use the pipeline to generate text.
try:
# The pipeline returns a list of dictionaries; we extract the generated text.
response = pipe(full_prompt, max_new_tokens=100)
return response[0]['generated_text']
except Exception as e:
return f"Error: {e}"
# Define the Gradio interface.
with gr.Blocks() as demo:
gr.Markdown("# Lightweight LLM Demo")
gr.Markdown("Enter text and select a prompt to generate an AI response.")
with gr.Row():
with gr.Column(scale=1):
# Textbox for user input
user_input = gr.Textbox(
label="Your Input Text",
placeholder="Type here...",
lines=5
)
# Dropdown to select a prompt template
prompt_template = gr.Dropdown(
label="Choose a Prompt Template",
choices=[
"Summarize this: {user_input}",
"Answer the following question: {user_input}",
"Rewrite this text to be more formal: {user_input}"
],
value="Summarize this: {user_input}"
)
# Button to trigger the generation
generate_button = gr.Button("Generate")
with gr.Column(scale=2):
# Textbox to display the output
output_text = gr.Textbox(
label="Generated Output",
lines=10
)
# Define the action for the button click
generate_button.click(
fn=generate_text,
inputs=[user_input, prompt_template],
outputs=output_text
)
demo.launch() |