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

# Check if model is extracted; if not, extract it
if not os.path.exists("best_model_subject"):
    shutil.unpack_archive("best_model_subject.zip", "best_model_subject")

# Load the saved model and tokenizer
model = AutoModelForSequenceClassification.from_pretrained("best_model_subject")
tokenizer = AutoTokenizer.from_pretrained("best_model_subject")

# Ensure the model is in evaluation mode
model.eval()

# Define the prediction function
def predict(Text):
    # Tokenize the input text
    inputs = tokenizer(Text, return_tensors="pt", max_length=512, truncation=True, padding=True)
    
    # Perform inference
    with torch.no_grad():
        logits = model(**inputs).logits
    
    # Get predicted label and confidence scores
    probs = torch.nn.functional.softmax(logits, dim=1)
    _, predicted_label = torch.max(logits, dim=1)
    
    # Map the predicted label to a human-readable class name
    class_names = ['Nature', 'Life-oriented', 'Rhyme', 'Allegorical', 'Love', 'Patriotic', 'Separation', 'Humanist', 'Religious']  
    predicted_class = class_names[predicted_label.item()]
    
    # Convert confidence scores to percentage with 2 decimal places
    probs_percentage = [f"{p * 100:.2f}%" for p in probs.tolist()[0]]
    
    # Return the predicted class and formatted confidence scores
    return predicted_class, str(probs_percentage)

# Create the Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(lines=2, placeholder="Enter your text"),
    outputs=[
        gr.Textbox(label="Predicted Class"),
        gr.Textbox(label="Confidence")
    ],
    title="Subject-based Classification",
    description="Classify poem into predefined categories."
)

# Launch the interface
iface.launch()