get_related_words.py 1.0 KB

12345678910111213141516171819202122232425262728
  1. from fastapi import APIRouter
  2. from gensim.models import KeyedVectors
  3. from loguru import logger
  4. import traceback
  5. router = APIRouter(tags=['related_words_server'])
  6. wv = KeyedVectors.load("/data/pythonProject/video_to_word/models/word2vec.wordvectors", mmap='r')
  7. @logger.catch()
  8. @router.post("/related_words/", description="获取相关词", summary="获取相关词")
  9. def get_related_words(word: str):
  10. try:
  11. similar_word_lst = wv.most_similar(word, topn=10)
  12. similar_word_lst = [item[0] for item in similar_word_lst if item[1] > 0.8]
  13. response = {'code': 0,
  14. "success": True,
  15. "result": similar_word_lst}
  16. logger.info(f"request body: {word}, response body: {response}")
  17. return response
  18. except:
  19. response = {"code": -1,
  20. "message": traceback.format_exc(),
  21. "success": False,
  22. "result": None}
  23. logger.error(f"request body: {word}, response body: {response}")
  24. return response