get_related_words.py 1.1 KB

12345678910111213141516171819202122232425262728293031323334
  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. response = {'code': 0,
  11. "success": True,
  12. "result": []}
  13. try:
  14. if word in wv:
  15. similar_word_lst = wv.most_similar(word, topn=10)
  16. similar_word_lst = [item[0] for item in similar_word_lst if item[1] > 0.73]
  17. response['result'] = similar_word_lst
  18. logger.info(f"request body: {word}, response body: {response}")
  19. return response
  20. except:
  21. response = {"code": -1,
  22. "message": traceback.format_exc(),
  23. "success": False,
  24. "result": []}
  25. logger.error(f"request body: {word}, response body: {response}")
  26. return response
  27. if __name__ == '__main__':
  28. get_related_words('社交')