import gradio as gr

import torch
from transformers import BertForSequenceClassification, BertTokenizer

# load model
tokenizer = BertTokenizer.from_pretrained("uget/sexual_content_dection")
model = BertForSequenceClassification.from_pretrained("uget/sexual_content_dection")

def predict(text):
    encoding = tokenizer(text, return_tensors="pt")
    encoding = {k: v.to(model.device) for k,v in encoding.items()}

    outputs = model(**encoding)
    probs = torch.sigmoid(outputs.logits)
    
    predictions = torch.argmax(probs, dim=-1)
    label_map = {0: "None", 1: "Sexual"}
    predicted_label = label_map[predictions.item()]
    print(f"Predictions:{predictions.item()}, Label:{predicted_label}")
    return {"predictions": predictions.item(), "label": predicted_label}



demo = gr.Interface(fn=predict, 
                    inputs="text", 
                    outputs="text",
                    examples=[["Tiffany Doll - Wine Makes Me Anal (31.03.2018)_1080p.mp4","{'predictions': 1, 'label': 'Sexual'}"],
                              ["DVAJ-548_CH_SD","{'predictions': 1, 'label': 'Sexual'}"],
                              ["MILK-217-UNCENSORED-LEAKピタコス Gカップ痴女 完全着衣で濃密5PLAY 椿りか 580 2.TS","{'predictions': 1, 'label': 'Sexual'}"],],
                    title="Sexual Content Detection",
                    description="Detects sexual content in text, <a href='https://ko-fi.com/ugetai' target='_blank'>Buy me a cup of coffee</a>.",
                    
                    )
demo.launch(share=True)