import gradio as gr
from transformers import pipeline


device = "cuda:0"

model = pipeline("text-classification", "iknow-lab/ko-flan-zero-v0-0731", device=device)
model.tokenizer.truncation_side = 'left'

def inference(input, instruction, labels):
    instruction = f"{input} [SEP] {instruction}"
    inputs = model.tokenizer([instruction] * len(labels), labels, truncation=True, padding=True, return_tensors="pt").to(device)
    
    scores = model.model(**inputs).logits.squeeze(1).softmax(-1).tolist()
    output = dict(zip(labels, scores))

    print(instruction)
    print(output)
    return output


def greet(content, instruction, labels):
    labels = labels.split(",")
    output = inference(content, instruction, labels)
    return output

content = gr.TextArea(label="입력 내용")
instruction = gr.Textbox(label="지시문")
labels = gr.Textbox(label="라벨(쉼표로 구분)")

examples = [
    ["예전에는 주말마다 극장에 놀러갔는데 요새는 좀 안가는 편이에요", "댓글 주제를 분류하세요", "영화,드라마,게임,소설"],
    ["인천발 KTX와 관련한 송도역 복합환승센터가 사실상 무산, 단순 철도·버스 위주 환승시설로 만들어진다. 이 때문에 인천시의 인천발 KTX 기점에 앵커시설인 복합환승센터를 통한 인근 지역 경제 활성화를 이뤄낸다는 계획의 차질이 불가피하다.", "경제에 긍정적인 뉴스인가요?", "예,아니요"],
    ["마지막에는 k팝 공연보고 좋은 추억 남았으면 좋겠네요","욕설이 포함되어있나요?", "욕설이 있습니다,욕설이 없습니다"],
]
gr.Interface(fn=greet,
             inputs=[content, instruction, labels],
             outputs=gr.Label(),
             examples=examples).launch() # server_name="0.0.0.0",server_port=7860)