# prompt: write a gradio app to infer the labels from the model we previously trained

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

# Load the fine-tuned model and tokenizer
checkpoint = "25b3nk/ollama-issues-classifier"  # Replace with the actual path to your checkpoint directory
model = AutoModelForSequenceClassification.from_pretrained(checkpoint)
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

# Move the model to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Function to perform inference
def predict(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
    outputs = model(**inputs)
    logits = outputs.logits
    probabilities = torch.sigmoid(logits)  # Use sigmoid for multi-label classification

    # Get predicted labels based on a threshold (e.g., 0.5)
    predicted_labels = (probabilities > 0.5).nonzero()[:, 1].tolist()

    # Map label IDs back to label names
    predicted_labels_names = [model.config.id2label[label_id] for label_id in predicted_labels]

    return predicted_labels_names


# Create the Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(lines=5, placeholder="Enter the issue text here..."),
    outputs=gr.Label(num_top_classes=len(model.config.id2label)),  # Display predicted labels
    title="Issue Label Prediction",
    description="Enter an issue description to predict its labels.",
)

iface.launch(debug=True)