import gradio as gr
import os
import requests

hf_token = os.getenv("HF_TOKEN")

api_url = "https://api-inference.huggingface.co/models/SARAL-Influencer-OS/comment_analyser"

headers = {"Authorization": f"Bearer {hf_token}"} if hf_token else {}

def predict_comment(comment):
    data = {"inputs": comment}
    response = requests.post(api_url, headers=headers, json=data)
    response.raise_for_status()  # Will raise an error for bad status codes
    output = response.json()

    
    label_map = {
        "LABEL_0": "Buying Intent",
        "LABEL_1": "Neutral Intent",
        "LABEL_2": "Negative Intent",
    }
    
    remapped_output = {label_map[item['label']]: item['score'] for item in output[0]}
    
    return remapped_output


iface = gr.Interface(
    fn=predict_comment,
    inputs="text",
    outputs="label"
)

iface.launch()