import gradio as gr
import onnxruntime as rt
from transformers import AutoTokenizer
import torch
import json

# Initialize the tokenizer
tokenizer = AutoTokenizer.from_pretrained("distilroberta-base")

# Load genre types from a JSON file
try:
    with open("genre_types_encoded.json", "r") as fp:
        encode_genre_types = json.load(fp)
except FileNotFoundError:
    print("Error: 'genre_types_encoded.json' not found. Make sure the file exists.")
    exit(1)

# Extract genres from the loaded data
genres = list(encode_genre_types.keys())

# Load the ONNX inference session
try:
    inf_session = rt.InferenceSession('udemy-classifier-quantized.onnx')
    input_name = inf_session.get_inputs()[0].name
    output_name = inf_session.get_outputs()[0].name
except FileNotFoundError:
    print("Error: 'udemy-classifier-quantized.onnx' not found. Make sure the file exists.")
    exit(1)

# Define the function for classifying courses' genres
def classify_courses_genre(description):
    input_ids = tokenizer(description, truncation=True, padding=True, return_tensors="pt")['input_ids'][:,:512]
    logits = inf_session.run([output_name], {input_name: input_ids.cpu().numpy()})[0]
    logits = torch.FloatTensor(logits)
    probs = torch.sigmoid(logits)[0]
    return dict(zip(genres, map(float, probs)))

# Create the Gradio interface
iface = gr.Interface(fn=classify_courses_genre, inputs="text", outputs=gr.components.Label(num_top_classes=5))

# Launch the Gradio interface
iface.launch(inline=False)