from fastapi import APIRouter
from gensim.models import KeyedVectors
from loguru import logger
import traceback

router = APIRouter(tags=['related_words_server'])

wv = KeyedVectors.load("/data/pythonProject/video_to_word/models/word2vec.wordvectors", mmap='r')


@logger.catch()
@router.post("/related_words/", description="获取相关词/推荐词", summary="获取相关词/推荐词")
def get_related_words(word: str):
    response = {'code': 0,
                "success": True,
                "result": []}
    try:
        if word in wv:
            similar_word_lst = wv.most_similar(word, topn=10)
            similar_word_lst = [item[0] for item in similar_word_lst if item[1] > 0.73]
            response['result'] = similar_word_lst
        logger.info(f"request body: {word}, response body: {response}")
        return response
    except:
        response = {"code": -1,
                    "message": traceback.format_exc(),
                    "success": False,
                    "result": []}
        logger.error(f"request body: {word}, response body: {response}")
        return response


if __name__ == '__main__':
    get_related_words('社交')
