import gradio as gr
from transformers import pipeline

# Initialize the model and pipeline
generate_text = pipeline(
    model="Stevross/Astrid-3B",
    task="text-generation",
    trust_remote_code=True,
    device=1  # 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="Stevross/Astrid-3B",
    description="An interactive demo of the Stevross/Astrid-3B model from Hugging Face."
)
 
interface.launch(share=True)