main.py

"""
This is main.py
"""
from fastapi import FastAPI, File, UploadFile, Form, Request
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, RedirectResponse, JSONResponse
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel
from typing import List

class AppData:
    def __init__(self):
        self.file_content = ""
        self.name_list = []
        self.place = []
        self.times = []
        self.name_dic = {}
        self.end_output = []


class ItemListRequest(BaseModel):
    nameList: List[str]


app_data = AppData()

# 설정
app = FastAPI()
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")


@app.get("/", response_class=HTMLResponse)
async def page_home(request: Request):
    """INDEX.HTML 화면"""
    return templates.TemplateResponse("index.html", {"request": request})


@app.get("/put.html", response_class=HTMLResponse)
async def page_put(request: Request):
    """PUT.HTML 화면"""
    return templates.TemplateResponse("put.html", {"request": request})


@app.get("/confirm.html", response_class=HTMLResponse)
async def page_confirm(request: Request):
    """confirm.HTML 화면"""
    return templates.TemplateResponse("confirm.html",{
        "request": request, "file_content": app_data.file_content})


@app.get("/result.html", response_class=HTMLResponse)
async def page_result(request: Request):
    """result.HTML 화면"""
    return templates.TemplateResponse("result.html", {"request": request})


@app.get("/user.html", response_class=HTMLResponse)
async def page_user(request: Request):
    """user.HTML 화면"""
    return templates.TemplateResponse("user.html", {"request": request})


@app.get("/final.html", response_class=HTMLResponse)
async def page_final(request: Request):
    """final.HTML 화면"""
    return templates.TemplateResponse("final.html", {"request": request,
                                                     "output": app_data.end_output,
                                                     "place": app_data.place,
                                                     "time": app_data.times})


@app.post("/upload", response_class=HTMLResponse)
async def upload_file(file: UploadFile = File(...)):
    """파일 업로드 및 저장"""
    with open("uploads/" + file.filename, "wb") as f:
        f.write(file.file.read())

    with open("uploads/" + file.filename, "r", encoding="utf-8") as f:
        app_data.file_content = f.read()

    return RedirectResponse(url="/put.html")


@app.post("/ners", response_class=JSONResponse)
async def ner_file():
    """저장된 파일에 NER 작업을 해서 화자랑 장소를 구분"""
    from utils.load_model import load_ner
    from utils.input_process import make_ner_input
    from utils.ner_utils import make_name_list, show_name_list, combine_similar_names

    content = app_data.file_content
    _, ner_checkpoint = load_ner()

    contents = make_ner_input(content)
    name_list, times, places = make_name_list(contents, ner_checkpoint)
    name_dic = show_name_list(name_list)
    similar_name = combine_similar_names(name_dic)
    result_list = [', '.join(names) for names, _ in similar_name.items()]
    app_data.place  = ' '.join(places)
    app_data.times = ' '.join(times)

    
    # JSONResponse로 응답
    return JSONResponse(content={"itemList": result_list})


@app.post("/kcsn", response_class=JSONResponse)
async def kcsn_file(request_data: ItemListRequest):
    """사용자가 올려준 파일에 대해서 KCSN 모델 동작"""
    import torch
    from utils.fs_utils import get_alias2id, find_speak, making_script
    from utils.input_process import make_instance_list, input_data_loader
    from utils.train_model import KCSN
    from utils.ner_utils import convert_name2codename, convert_codename2name

    content = app_data.file_content
    name_list = request_data.nameList
    name_dic = {}

    for idx, name in enumerate(name_list):
        name_dic[f'&C{idx:02d}&'] = name.split(', ')

    content_re = convert_name2codename(name_dic, content)

    # checkpoint = torch.load('./model/final.pth')
    # model = checkpoint['model']
    # model.to('cpu')
    # tokenizer = checkpoint['tokenizer']

    from utils.arguments import get_train_args
    from transformers import AutoTokenizer

    args = get_train_args()
    path ='model/model.ckpt'
    model = KCSN(args)
    model.to('cpu')

    checkpoint = torch.load(path)
    tokenizer = AutoTokenizer.from_pretrained(args.bert_pretrained_dir)
    model.load_state_dict(checkpoint['model'])

    check_name = 'data/name.txt'
    alias2id = get_alias2id(check_name)
    instances, instance_num = make_instance_list(content_re)
    inputs = input_data_loader(instances, alias2id)
    output = find_speak(model, inputs, tokenizer, alias2id)
    outputs = convert_codename2name(name_dic, output)
    app_data.end_output = making_script(content, outputs, instance_num)


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="127.0.0.1", port=8000)