import gradio as gr
import subprocess
import numpy as np

# Install the required libraries if not already installed
required_libraries = ["transformers", "torch"]

for lib in required_libraries:
    try:
        __import__(lib)
    except ImportError:
        subprocess.run(["pip", "install", lib])

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Load your Hugging Face model.
model_name = "Norah-K/PDPL_CAMeLBERT"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

def privacy_policy_classification(text):
    # Tokenize the input text using the Hugging Face tokenizer.
    inputs = tokenizer(text, truncation=True, padding=True, max_length=512, return_tensors="pt")

    # Classify the privacy policy using the loaded model.
    with torch.no_grad():
        classification = model(**inputs)

    # Assuming your model's output is a classification label.
    # Calculate softmax to get probabilities
    probabilities = torch.softmax(classification.logits, dim=1).squeeze().tolist()

    # You can define a list of fixed labels for your classes here.
    fixed_labels = [
        "موافقة المستخدم",
        "بيانات التواصل أو تحت السن القانوني",
         "جمع ومعالجة البيانات",
         "الاحتفاظ بالبيانات",
         "حماية البيانات",
         "مشاركة البيانات",
         "حقوق المستخدم",
        "الاعلانات",
         "تنبيهات الاختراق",
        "المسؤوليات", 
    ]

    # Sort labels and probabilities by probability in descending order
    sorted_labels = [label for _, label in sorted(zip(probabilities, fixed_labels), reverse=True)]
    sorted_probabilities = sorted(probabilities, reverse=True)

    # Create a result string with labels and their corresponding probabilities, arranged by highest probability
    result = "\n".join([f"{probability * 100:.2f} % :{label}" for label, probability in zip(sorted_labels, sorted_probabilities)])
    
    return result

# Define a Gradio interface with the "gradio/soft" theme
iface = gr.Interface(
    fn=privacy_policy_classification,
    inputs="text",
    outputs="text",  # Output as plain text
    title="<b>PDPL Compliance Detector</b>",
    theme="gradio/soft",  # Set the theme to "gradio/soft"
)

# Launch the Gradio interface
iface.launch()