import gradio as gr
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# Load the model and tokenizer
model = AutoModelForSequenceClassification.from_pretrained("smeintadmin/image_intents")
tokenizer = AutoTokenizer.from_pretrained("smeintadmin/image_intents")

# Define prediction function
def predict(text):

    # Tokenize the text
    inputs = tokenizer(text, truncation=True, max_length=512, padding='max_length', return_tensors='pt')

    # Make a prediction
    outputs = model(**inputs)
    logits = outputs.logits
    predicted_class = torch.argmax(logits).item()

    # Return the label
    return "Image Intent" if predicted_class == 1 else "Non-Image Intent"

# Function to display user history
def history_display(history):
    return "\n".join(f"{data['input']} -> {data['output']}" for data in history)

# Custom layout using HTML tags and CSS
description_html = "<h4>Description</h4><p>This app classifies whether the input text represents an image intent or not.</p>"
instructions_html = "<h2>Instructions</h2><p>Enter a text query and click <b>Submit</b> to see if the intent is related to an image or not.</p>"

# Combine the HTML components to form the final layout
iface = gr.Interface(
    fn=predict,
    inputs="text",
    outputs="text",
    title="Image Intent Classifier",
    description=description_html,
    article=instructions_html,  # Use 'article' instead of 'inputs' for the instructions
    examples=[
        ["I want to see the picture."],
        ["Find me some recipes for dinner."],
        ["Can you tell me the latest news?"],
        ["Show me cute cat pictures!"],
    ],
    theme="huggingface",  # Dark mode theme
    allow_flagging="never",  # Disable flagging option (use 'auto' if you want to enable it)
    history=history_display,  # Function to display user history
)

# Launch the interface
iface.launch()