import pip
pip.main(['install', 'torch'])
pip.main(['install', 'transformers'])

import re

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

def load_model(model_name):
    # model
    model = AutoModelForSequenceClassification.from_pretrained(model_name)
    # tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    return model, tokenizer


def inference(prompt_inputs):

    model_name = "Unggi/feedback_prize_kor"

    model, tokenizer = load_model(
        model_name = model_name
    )

    # preprocessing
    prompt_inputs = prompt_inputs.replace('\n', ' ')

    # prompt 구두점 단위로 분리하기
    prompt_list = re.split('[.?!]+', prompt_inputs)

    class_id_list = []

    for prompt in prompt_list:
        inputs = tokenizer(
            prompt, 
            return_tensors="pt"
            )
    
        with torch.no_grad():
            logits = model(**inputs).logits
    
        predicted_class_id = logits.argmax().item()
        class_id = model.config.id2label[predicted_class_id]

        class_id_list.append(class_id)

    outputs = []

    for p, c_id in zip(prompt_list, class_id_list):

        outputs.append(p + '\t' + '=>' + '\t' + '<' + c_id + '>')

    outputs = '\n'.join(outputs)

    return outputs

demo = gr.Interface(
    fn=inference, 
    inputs="text", 
    outputs="text", #return 값
    examples=[
        "예전부터 여성의 사회참여운동이나 시위를 지속적이며 힘 있는 행동으로 확대해나갔으며 현재까지 계속 양성평등의 중요성이 주목받고 있다. 앞에서 말했듯이 아이 이름의 성을 아버지의 성으로 짓던 것을 이제 부모 모두를 바탕으로 이름을 짓는 양계제도를 도입하기도 하고 직장에서의 남녀 모두 승진과 임금 등을 동등하게 분배하며 남녀가 정해져 있던 직업들이 성차별의 경계가 무너짐과 동시에 누구나 할 수 있도록 하는 것 이외에 많은 양성평등의 예들이 도래되고 있다. 요약해서 성차별이 개개인을 구속하고 집단의 걸림돌이 되어 사람마다 자신의 주장과 행동에 대한 제한이 불가피하며 사회의 발전을 저해하는 것뿐만 아니라 선진국 또는 바람직한 국가로 나아가는데 방해가 된다는 것인데 이를 극복하기 위해서는 성차별에 대한 생각을 떨쳐버리고 양성평등이란 선진적 시민의식을 갖춤으로 선진국을 위해 국가의 노력도 필요하지만, 무엇보다 개인이 성차별에 대한 관념을 잊고 양성평등을 위한 적극적인 노력이 요구되어야 한다. 그로 인해 국가와 개인에 의한 바람직한 사회가 형성된다는 것이다."
    ]
    ).launch() # launch(share=True)를 설정하면 외부에서 접속 가능한 링크가 생성됨

demo.launch()