import gradio as gr
from transformers import pipeline

# Initialize the text-generation pipeline with your model
pipe = pipeline("text-generation", model="cognitivecomputations/dolphin-2.8-mistral-7b-v02")

# Define a function to generate a response based on user input
def generate_response(user_input):
    # The pipeline expects a list of dictionaries with keys 'role' and 'content'
    messages = [{"role": "user", "content": user_input}]
    
    # Generate the response using the pipeline
    response = pipe(messages, max_length=100)[0]['generated_text']
    
    return response

# Create a Gradio interface
iface = gr.Interface(
    fn=generate_response,
    inputs="text",
    outputs="text",
    title="AI Chatbot",
    description="A chatbot interface using a text-generation model.",
    theme="compact",
    layout="vertical",
)

# Launch the Gradio app
iface.launch()