import gradio as gr
import torch
from model_def import RNN
from name_change import hangul_to_roman, lineToTensor
import torch.nn.functional as F


# 모델 로드
RNN(input_size=57, hidden_size=128, output_size=2)
rnn_load = torch.load('rnn_08140037.pt')
rnn_load.eval()

# 주어진 라인의 출력 반환
all_categories=['여자','남자']

def evaluate(line_tensor):
    hidden = rnn_load.initHidden()

    for i in range(line_tensor.size()[0]):
        output, hidden = rnn_load(line_tensor[i], hidden)

    return output

def predict(input_line, n_predictions=2):
    print('\n> %s' % input_line)
    input_line=hangul_to_roman(input_line)
    with torch.no_grad():
        output = evaluate(lineToTensor(input_line))

        probabilities = F.softmax(output, dim=1)
        topv, topi = torch.topk(probabilities, n_predictions)
        predictions = [(round(topv[0][i].item(), 2), all_categories[topi[0][i].item()]) for i in range(n_predictions)]
    return predictions

def name_classifier(name):
    result=predict(name)
    print(result)
    return {result[0][1]: result[0][0], result[1][1]: result[1][0]}


demo = gr.Interface(
    fn=name_classifier, 
    inputs="text", 
    outputs="label",
    title="한국이름 성별 예측 모델",  
    description="이 모델은 입력된 이름을 기반으로 성별을 예측합니다. 성을 제외한 이름을 입력하고 예측된 성별과 그 확률을 확인하세요.",  
    examples=[["하은"], ["현배"], ["혜윤"]] 
)



demo.launch()