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

# # Set device to CPU since GPU quantization is unavailable
# device = 'cpu'

# # Set up 8-bit quantization with BitsAndBytesConfig
# quantization_config = BitsAndBytesConfig(
#     load_in_8bit=True,  # Enable 8-bit quantization
#     llm_int8_enable_fp32_cpu_offload=True  # Use CPU for 8-bit quantization operations
# )

# # Load the model with quantization configuration
# model_name = "Rahmat82/DistilBERT-finetuned-on-emotion"
# model = AutoModelForSequenceClassification.from_pretrained(
#     model_name,
#     quantization_config=quantization_config,
#     device_map={"": device}  # Ensures everything runs on CPU
# )
# tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)

# def predict(query: str) -> dict:
#     inputs = tokenizer(query, return_tensors='pt')
#     inputs = {k: v.to(device) for k, v in inputs.items()}  # Ensure inputs are on CPU
#     outputs = model(**inputs)
#     outputs = torch.sigmoid(outputs.logits)
#     outputs = outputs.detach().cpu().numpy()
    
#     # Define label to ID mapping
#     label2ids = {
#         "sadness": 0,
#         "joy": 1,
#         "love": 2,
#         "anger": 3,
#         "fear": 4,
#         "surprise": 5,
#     }
#     for i, k in enumerate(label2ids.keys()):
#         label2ids[k] = outputs[0][i]
#     label2ids = {k: float(v) for k, v in sorted(label2ids.items(), key=lambda item: item[1], reverse=True)}
#     return label2ids

# # Gradio interface setup
# demo = gr.Interface(
#     theme=gr.themes.Soft(),
#     title="RHM Emotion Classifier 😊",
#     description="Beyond Words: Capturing the Essence of Emotion in Text<h3>On CPU with 8-bit quantization</h3>",
#     fn=predict,
#     inputs=gr.components.Textbox(label='Write your text here', lines=3),
#     outputs=gr.components.Label(label='Predictions', num_top_classes=6),
#     allow_flagging='never',
#     examples=[
#         ["The gentle touch of your hand on mine is a silent promise that echoes through the corridors of my heart."],
#         ["The rain mirrored the tears I couldn't stop, each drop a tiny echo of the ache in my heart. The world seemed muted, colors drained, and a heavy weight settled upon my soul."],
#         ["Walking through the dusty attic, I stumbled upon a hidden door. With a mix of trepidation and excitement, I pushed it open, expecting cobwebs and forgotten junk. Instead, a flood of sunlight revealed a secret garden, blooming with vibrant flowers and buzzing with life. My jaw dropped in pure astonishment."],
#     ]
# )

# demo.launch(share=True)




import gradio as gr
from transformers import pipeline, AutoTokenizer
from optimum.onnxruntime import ORTModelForSequenceClassification
import torch

device = 'cuda' if torch.cuda.is_available() else 'cpu'

model_name = "Rahmat82/DistilBERT-finetuned-on-emotion"
model = ORTModelForSequenceClassification.from_pretrained(model_name, export=True)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
model.to(device)

def predict(query: str) -> dict:
    inputs = tokenizer(query, return_tensors='pt')
    inputs.to(device)
    outputs = model(**inputs)
    outputs = torch.sigmoid(outputs.logits)
    outputs = outputs.detach().cpu().numpy()
    
    label2ids = {
                "sadness": 0,
                "joy": 1,
                "love": 2,
                "anger": 3,
                "fear": 4,
                "surprise": 5,
                }
    for i, k in enumerate(label2ids.keys()):
        label2ids[k] = outputs[0][i]
    label2ids = {k: float(v) for k, v in sorted(label2ids.items(), key=lambda item: item[1], reverse=True)}
    return label2ids

demo = gr.Interface(
    theme = gr.themes.Soft(),
    title = "RHM Emotion Classifier 😊",
    description = "Beyond Words: Capturing the Essence of Emotion in Text<h3>On GPU it is much faster 🚀</h3>",
    fn = predict,
    inputs = gr.components.Textbox(label='Write your text here', lines=3),
    outputs = gr.components.Label(label='Predictions', num_top_classes=6),
    allow_flagging = 'never',
    examples = [
        ["The gentle touch of your hand on mine is a silent promise that echoes through the corridors of my heart."],
        ["The rain mirrored the tears I couldn't stop, each drop a tiny echo of the ache in my heart. The world seemed muted, colors drained, and a heavy weight settled upon my soul."],
        ["Walking through the dusty attic, I stumbled upon a hidden door. With a mix of trepidation and excitement, I pushed it open, expecting cobwebs and forgotten junk. Instead, a flood of sunlight revealed a secret garden, blooming with vibrant flowers and buzzing with life. My jaw dropped in pure astonishment."],
              ]
)

demo.launch(share=True)