{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "### 获得单例文本预测结果"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "from config import *\n",
    "from utils import *\n",
    "\n",
    "q_std_list, fst_q_std_sentence_embeddings, q_all, fst_q_all_sentence_embeddings_dict = read_q_std_q_corpus(q_std_file, fst_q_std_vectors_file, q_corpus_file, fst_q_corpus_vectors_file)\n",
    "_, sec_q_std_sentence_embeddings, _, sec_q_all_sentence_embeddings_dict = read_q_std_q_corpus(q_std_file, sec_q_std_vectors_file, q_corpus_file, sec_q_corpus_vectors_file)\n",
    "\n",
    "def get_fst_topK(text, K=10):\n",
    "    text_embedding = model1.encode([text])[0].numpy()\n",
    "    sims_with_std = np.array(cos_sim4matrix_2(text_embedding, fst_q_std_sentence_embeddings))\n",
    "    sort_idx = np.argsort(-sims_with_std)[:K]\n",
    "    sims_q_sort = [q_std_list[idx] for idx in sort_idx]\n",
    "    sims_values = [sims_with_std[idx] for idx in sort_idx]\n",
    "    result = list(zip(sims_q_sort, sims_values))\n",
    "    return (result)\n",
    "\n",
    "def get_sec_topK(self, text, K=20):\n",
    "    text_embedding = self.model.encode([text])[0]\n",
    "    sims_with_std = np.array(self.cos_sim4matrix_2(text_embedding, self.std_sentence_embeddings))\n",
    "    sort_idx = np.argsort(-sims_with_std)[:K]\n",
    "    sims_q_sort = [self.q_std_list[idx] for idx in sort_idx]\n",
    "    sims_values = [sims_with_std[idx] for idx in sort_idx]\n",
    "    result = list(zip(sims_q_sort, sims_values))\n",
    "    return (result)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import numpy as np\n",
    "from config import *\n",
    "from utils import *\n",
    "from task_sentence_embedding_FinanceFAQ_step1_1 import model as model1\n",
    "from task_sentence_embedding_FinanceFAQ_step2_1 import model as model2\n",
    "\n",
    "text = input()\n",
    "\n",
    "# 第一阶段召回\n",
    "result_first = get_fst_topK(text=text)\n",
    "print('第一阶段\\n', result_first[0:20])\n",
    "first_intents = [v[0] for v in result_first]\n",
    "\n",
    "# 第二阶段召回\n",
    "a_texts_embeddings_2 = np.array(model2.encode([text]))\n",
    "b_texts_embeddings_2 = np.array([sec_q_all_sentence_embeddings_dict[v] for v in first_intents])\n",
    "sims_with_std = cos_sim4matrix_2(a_texts_embeddings_2, b_texts_embeddings_2).reshape(-1)\n",
    "sort_idx = np.argsort(-sims_with_std).tolist()\n",
    "intents_sort = [first_intents[idx] for idx in sort_idx]\n",
    "sims_values = [sims_with_std[idx] for idx in sort_idx]\n",
    "result_second = list(zip(intents_sort, sims_values))\n",
    "print('第二阶段\\n', result_second[0:20])"
   ],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "orig_nbformat": 4,
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}