import torch
import torch.nn as nn
import torch.nn.functional as F
import os
    
from model_def import TextClassifier
from mor import tokenize
import pickle
import gradio as gr
import subprocess




embedding_dim = 100
hidden_dim = 128
output_dim = 2
vocab_size=17391 
USE_CUDA = torch.cuda.is_available()
device = torch.device("cuda" if USE_CUDA else "cpu")
model_name='08221228'

model = TextClassifier(vocab_size, embedding_dim, hidden_dim, output_dim)


model.load_state_dict(torch.load('best_model_checkpoint'+model_name+'.pth',map_location=device))
model.to(device)

with open('word_to_index.pkl', 'rb') as f:
    word_to_index = pickle.load(f)




index_to_tag = {0 : '부정', 1 : '긍정'}
def predict(text, model, word_to_index, index_to_tag):
    # Set the model to evaluation mode
    model.eval()
    tokens= tokenize(text)
    
    token_indices = [word_to_index.get(token, 1) for token in tokens]

    input_tensor = torch.tensor([token_indices], dtype=torch.long).to(device)

    # Pass the input tensor through the model
    with torch.no_grad():
        logits = model(input_tensor)  # (1, output_dim)

    # Apply softmax to the logits
    probs = F.softmax(logits, dim=1)
    topv, topi = torch.topk(probs, 2)
    predictions = [(round(topv[0][i].item(), 2), index_to_tag[topi[0][i].item()]) for i in range(2)]

    # Get the predicted class index
    predicted_index = torch.argmax(logits, dim=1)

    # Convert the predicted index to its corresponding tag
    predicted_tag = index_to_tag[predicted_index.item()]

    return predictions



def name_classifier(test_input):
    result=predict(test_input, model, word_to_index, index_to_tag)
    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="영화 리뷰 감성 분석 LSTM 모델",  
    description="이 모델은 영화 리뷰 텍스트를 입력받아 감성 분석을 수행하여, 긍정적 또는 부정적인 감정을 예측합니다. LSTM 기반의 텍스트 분류 모델입니다. 이 모델은 위키독스의 [13-02 LSTM을 이용한 네이버 영화 리뷰 분류](https://wikidocs.net/217687)를 바탕으로 제작한 예제입니다.",
    examples=[["뭔가 맺음이 없는 느낌.."], [" 하츄핑과 로미의 사랑이야기...의외로 ost가 너무 좋아요! "]] 
)



demo.launch()