# import argparse
# import os
# from importlib import import_module

# import gradio as gr
# from tqdm import tqdm
# import models.TextCNN
# import torch
# import pickle as pkl
# from utils import build_dataset

# classes = ['金融类', '房地产类', '股票类', '教育类', '科技类', '社会类', '政治类', '体育类', '游戏类',
#            '娱乐类']

# MAX_VOCAB_SIZE = 10000  # 词表长度限制
# UNK, PAD = '<UNK>', '<PAD>'  # 未知字，padding符号


# def build_vocab(file_path, tokenizer, max_size, min_freq):
#     vocab_dic = {}
#     with open(file_path, 'r', encoding='UTF-8') as f:
#         for line in tqdm(f):
#             lin = line.strip()
#             if not lin:
#                 continue
#             content = lin.split('\t')[0]
#             for word in tokenizer(content):
#                 vocab_dic[word] = vocab_dic.get(word, 0) + 1
#         vocab_list = sorted([_ for _ in vocab_dic.items() if _[1] >= min_freq], key=lambda x: x[1], reverse=True)[
#                      :max_size]
#         vocab_dic = {word_count[0]: idx for idx, word_count in enumerate(vocab_list)}
#         vocab_dic.update({UNK: len(vocab_dic), PAD: len(vocab_dic) + 1})
#     return vocab_dic








# def greet(text):
#     parser = argparse.ArgumentParser(description='Chinese Text Classification')
#     parser.add_argument('--word', default=False, type=bool, help='True for word, False for char')
#     args = parser.parse_args()
#     model_name = 'TextCNN'
#     dataset = 'THUCNews'  # 数据集
#     embedding = 'embedding_SougouNews.npz'
#     x = import_module('models.' + model_name)

#     config = x.Config(dataset, embedding)
#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

#     model = models.TextCNN.Model(config)

#     # vocab, train_data, dev_data, test_data = build_dataset(config, args.word)
#     model.load_state_dict(torch.load('THUCNews/saved_dict/TextCNN.ckpt', map_location=torch.device('cpu')))
#     model.to(device)
#     model.eval()

#     tokenizer = lambda x: [y for y in x]  # char-level
#     if os.path.exists(config.vocab_path):
#         vocab = pkl.load(open(config.vocab_path, 'rb'))
#     else:
#         vocab = build_vocab(config.train_path, tokenizer=tokenizer, max_size=MAX_VOCAB_SIZE, min_freq=1)
#         pkl.dump(vocab, open(config.vocab_path, 'wb'))
#     # print(f"Vocab size: {len(vocab)}")

#     # content='时评：“国学小天才”录取缘何少佳话'
#     content = text

#     words_line = []
#     token = tokenizer(content)
#     seq_len = len(token)
#     pad_size = 32
#     contents = []

#     if pad_size:
#         if len(token) < pad_size:
#             token.extend([PAD] * (pad_size - len(token)))
#         else:
#             token = token[:pad_size]
#             seq_len = pad_size
#     # word to id
#     for word in token:
#         words_line.append(vocab.get(word, vocab.get(UNK)))

#     contents.append((words_line, seq_len))
#     # print(words_line)
#     # input = torch.LongTensor(words_line).unsqueeze(1).to(device)  # convert words_line to LongTensor and add batch dimension
#     x = torch.LongTensor([_[0] for _ in contents]).to(device)

#     # pad前的长度(超过pad_size的设为pad_size)
#     seq_len = torch.LongTensor([_[1] for _ in contents]).to(device)
#     input = (x, seq_len)
#     # print(input)
#     with torch.no_grad():
#         output = model(input)
#         predic = torch.max(output.data, 1)[1].cpu().numpy()
#     # print(predic)
#     # print('类别为：{}'.format(classes[predic[0]]))
#     return classes[predic[0]]





# demo = gr.Interface(fn=greet, inputs="text", outputs="text", title="text-classification app", 
#                      layout="vertical", description="This is a demo for text classification.")
# demo.launch()

#你可以使用 CSS 和 HTML 来自定义你的 Gradio 界面，以使其更具吸引力。以下是一个示例，其中使用了一些 CSS 样式和 HTML 标记来改进界面布局和风格：


import argparse
import os
from importlib import import_module

import gradio as gr
import models.TextCNN
import torch
import pickle as pkl
from tqdm import tqdm

classes = ['金融类', '房地产类', '股票类', '教育类', '科技类', '社会类', '政治类', '体育类', '游戏类', '娱乐类']

MAX_VOCAB_SIZE = 10000  # 词表长度限制
UNK, PAD = '<UNK>', '<PAD>'  # 未知字，padding符号


def build_vocab(file_path, tokenizer, max_size, min_freq):
    vocab_dic = {}
    with open(file_path, 'r', encoding='UTF-8') as f:
        for line in tqdm(f):
            lin = line.strip()
            if not lin:
                continue
            content = lin.split('\t')[0]
            for word in tokenizer(content):
                vocab_dic[word] = vocab_dic.get(word, 0) + 1
        vocab_list = sorted([_ for _ in vocab_dic.items() if _[1] >= min_freq], key=lambda x: x[1], reverse=True)[
                     :max_size]
        vocab_dic = {word_count[0]: idx for idx, word_count in enumerate(vocab_list)}
        vocab_dic.update({UNK: len(vocab_dic), PAD: len(vocab_dic) + 1})
    return vocab_dic


def greet(text):
    parser = argparse.ArgumentParser(description='Chinese Text Classification')
    parser.add_argument('--word', default=False, type=bool, help='True for word, False for char')
    args = parser.parse_args()
    model_name = 'TextCNN'
    dataset = 'THUCNews'  # 数据集
    embedding = 'embedding_SougouNews.npz'
    x = import_module('models.' + model_name)

    config = x.Config(dataset, embedding)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model = models.TextCNN.Model(config)

    # vocab, train_data, dev_data, test_data = build_dataset(config, args.word)
    model.load_state_dict(torch.load('THUCNews/saved_dict/TextCNN.ckpt', map_location=torch.device('cpu')))
    model.to(device)
    model.eval()

    tokenizer = lambda x: [y for y in x]  # char-level
    if os.path.exists(config.vocab_path):
        vocab = pkl.load(open(config.vocab_path, 'rb'))
    else:
        vocab = build_vocab(config.train_path, tokenizer=tokenizer, max_size=MAX_VOCAB_SIZE, min_freq=1)
        pkl.dump(vocab, open(config.vocab_path, 'wb'))
    # print(f"Vocab size: {len(vocab)}")

    # content='时评：“国学小天才”录取缘何少佳话'
    content = text

    words_line = []
    token = tokenizer(content)
    seq_len = len(token)
    pad_size = 32
    contents = []

    if pad_size:
        if len(token) < pad_size:
            token.extend([PAD] * (pad_size - len(token)))
        else:
            token = token[:pad_size]
            seq_len = pad_size
    # word to id
    for word in token:
        words_line.append(vocab.get(word, vocab.get(UNK)))

    contents.append((words_line, seq_len))
    # print(words_line)
    # input = torch.LongTensor(words_line).unsqueeze(1).to(device)  # convert words_line to LongTensor and add batch dimension
    x = torch.LongTensor([_[0] for _ in contents]).to(device)

    # pad前的长度(超过pad_size的设为pad_size)
    seq_len = torch.LongTensor([_[1] for _ in contents]).to(device)
    input = (x, seq_len)
    # print(input)
    with torch.no_grad():
        output = model(input)
        predic = torch.max(output.data, 1)[1].cpu().numpy()
    # print(predic)
    # print('类别为：{}'.format(classes[predic[0]]))
    return classes[predic[0]]

# 自定义样式和布局
css = """
body {
    background-color: #f8f8f8;
    font-family: Arial, sans-serif;
}

.container {
    max-width: 800px;
   margin: 0 auto;
    padding: 50px;
}

h1 {
    font-size: 36px;
    font-weight: bold;
    color: #333333;
    text-align: center;
    margin-bottom: 50px;
}

.gradio-interface {
    border: none;
    box-shadow: 0px 0px 10px rgba(0, 0, 0, 0.1);
    border-radius: 10px;
    overflow: hidden;
    margin-bottom: 50px;
}

.gradio-input {
    background-color: #ffffff;
    border: none;
    border-radius: 5px;
    padding: 15px;
    box-shadow: 0px 0px 10px rgba(0, 0, 0, 0.1);
    font-size: 18px;
    color: #333333;
    width: 100%;
    margin-bottom: 20px;
}

.gradio-output {
    background-color: #ffffff;
    border: none;
    border-radius: 5px;
    padding: 15px;
    box-shadow: 0px 0px 10px rgba(0, 0, 0, 0.1);
    font-size: 18px;
    color: #333333;
    width: 100%;
    margin-bottom: 20px;
}

.gradio-interface:hover {
    box-shadow: 0px 0px 20px rgba(0, 0, 0, 0.2);
}

.gradio-interface:focus-within {
    box-shadow: 0px 0px 20px rgba(0, 0, 0, 0.2);
}

.gradio-interface .input-group {
    margin-bottom: 20px;
}

.gradio-interface .input-label {
    font-size: 24px;
    font-weight: bold;
    color: #333333;
    margin-bottom: 10px;
}

.gradio-interface .input-description {
    font-size: 16px;
    color: #666666;
    margin-bottom: 20px;
}

.gradio-interface .output-label {
    font-size: 24px;
    font-weight: bold;
    color: #333333;
    margin-bottom: 10px;
}

.gradio-interface .output-description {
    font-size: 16px;
    color: #666666;
    margin-bottom: 20px;
}

.gradio-interface .input-group input[type="text"]::placeholder {
    color: #999999;
}

.gradio-button {
    background-color: #333333;
    color: #ffffff;
    border: none;
    border-radius: 5px;
    padding: 15px 30px;
    font-size: 18px;
    font-weight: bold;
    cursor: pointer;
    transition: background-color 0.2s ease;
}

.gradio-button:hover {
    background-color: #111111;
}
"""

html = """
<div class="container">
    <h1>Text Classification</h1>
    <div class="gradio-interface">
        <div class="input-group">
            <label class="input-label">Input Text</label>
            <div class="input-description">Enter the text that you want to classify:</div>
            <input type="text" class="gradio-input" placeholder="Enter your text here" id="input_text">
        </div>
        <button type="button" class="gradio-button" id="classify_button">Classify</button>
        <div class="output-group">
            <label class="output-label">Classification Result</label>
            <div class="output-description">The predicted class is:</div>
            <div class="gradio-output" id="output_text"></div>
        </div>
    </div>
</div>
"""

iface = gr.Interface(fn=greet, inputs="text", outputs="text", title="Text Classification App",
                     layout="unaligned", description="This is a demo for text classification.", css=css, 
                     allow_screenshot=False, allow_flagging=False, allow_share=False, allow_download=False,
                     examples=[["今天天气真好"], ["这个手机真不错"], ["新冠疫情对经济的影响"]])

iface.launch(html=html)