import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline

model_name = "SJ-Donald/kcbert-large-unsmile"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)

def classify_text(text):
    """
    주어진 텍스트가 부적절한지 여부를 판별합니다.
    
    Args:
    - text (str): 판별할 텍스트
    
    Returns:
    - result (str): '부적절한 내용이 포함되어 있습니다.' or '적절한 내용입니다.'
    """
    results = classifier(text)
    print(f"Debugging results: {results}")  # 결과 확인용 출력    
    
    for result in results:
        # 모델에 따라 라벨이 다를 수 있습니다.
        if result['label'] == '악플/욕설' and result['score'] > 0.5:
            return "악플/욕설입니다."
        elif result['label'] == '여성/가족' and result['score'] > 0.5:
            return "여성 혐오입니다"
        elif result['label'] == '남성' and result['score'] > 0.5:
            return "남성 혐오입니다"
        elif result['label'] == '인종/국적' and result['score'] > 0.5:
            return "인종/국적 혐오입니다"
        elif result['label'] == '연령' and result['score'] > 0.5:
            return "연령 혐오입니다"
        elif result['label'] == '지역' and result['score'] > 0.5:
            return "지역 혐오입니다"
        elif result['label'] == '종교' and result['score'] > 0.5:
            return "종교 혐오입니다"
        elif result['label'] == '기타 혐오' and result['score'] > 0.5:
            return "기타 혐오입니다"
    
    return "적절한 내용입니다."



demo = gr.Interface(fn=classify_text, inputs="textbox", title="부적절 문장 검출기", theme="soft", description="기준: 여성/가족, 남성, 성소수자, 인종/국적, 연령, 지역, 종교, 기타 혐오, 악플/욕설", outputs="textbox")
demo.launch()