from transformers import pipeline
import gradio as gr

# Initialize the pipeline
pipe = pipeline("text-generation", model="nvidia/OpenMath2-Llama3.1-8B")

# Define a function that takes user input and uses the pipeline to generate a response
def generate_response(messages):
    # Prepare the input format for text generation
    conversation = [{"role": "user", "content": messages}]
    response = pipe(conversation)[0]['generated_text']
    return response

# Set up the Gradio interface
iface = gr.Interface(
    fn=generate_response,  # Function to call
    inputs="text",  # Input type: text
    outputs="text",  # Output type: text
    title="Text Generation with Llama 3.1",
    description="Ask questions or have a conversation with Llama 3.1",
)

# Launch the interface
iface.launch()