import gradio as gr
from transformers import pipeline

# Initialize the model and pipeline
generate_text = pipeline(
    model="Stevross/Astrid-7B-Assistant-CPU",
    task="text-generation",
    device=0  # Use GPU if available, otherwise CPU
)

def generate_response(prompt):
    response = generate_text(prompt, max_length=150)
    return response[0]['generated_text']

# Define Gradio interface
interface = gr.Interface(
    fn=generate_response,
    inputs=gr.inputs.Textbox(lines=5, placeholder="Enter your prompt here..."),
    outputs=gr.outputs.Textbox(),
    live=True,
    title="Astrid-7B-Assistant-CPU Demo",
    description="An interactive demo of the Astrid-7B-Assistant-CPU model from Hugging Face."
)

interface.launch()