go00od's picture
Upload app.py with huggingface_hub
3138369 verified
raw
history blame
2.5 kB
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()