import gradio as gr
from transformers import pipeline

# Load the text-generation pipeline with Mistral model
pipe = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.3")

# Define the function to process user input
def classify_text(text):
    prompt = "Classify the following text into a category or topic:"
    input_text = f"{prompt}\n{text}"
    results = pipe(input_text, max_length=100, num_return_sequences=1)
    return results[0]['generated_text']

# Create Gradio interface
interface = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(lines=4, placeholder="Enter your text here..."),
    outputs=gr.Textbox(lines=4),
    title="Text Classification with Mistral",
    description="Enter some text to classify it into a category or topic using the Mistral-7B-Instruct-v0.3 model."
)

# Launch the app
if __name__ == "__main__":
    interface.launch()