{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d82abfc8-1e95-41f0-a9af-4946de3ad846",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ДХЛ. Красная шапочка и другие сказки\n",
      "Ослиная шкура\n",
      "Рождественское чудо мистера Туми\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import torch\n",
    "import numpy as np\n",
    "from transformers import AutoTokenizer, AutoModel\n",
    "import faiss\n",
    "\n",
    "# Загрузка модели и токенизатора BERT\n",
    "model_name = \"cointegrated/rubert-tiny2\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "model = AutoModel.from_pretrained(model_name)\n",
    "\n",
    "# Загрузка данных из CSV\n",
    "df = pd.read_csv('final_data.csv')\n",
    "\n",
    "# Максимальная длина текста\n",
    "MAX_LEN = 300\n",
    "\n",
    "# Функция для встраивания текста с использованием BERT\n",
    "def embed_bert_cls(text, model=model, tokenizer=tokenizer):\n",
    "    t = tokenizer(text,\n",
    "                  padding=True,\n",
    "                  truncation=True,\n",
    "                  return_tensors='pt',\n",
    "                  max_length=MAX_LEN)\n",
    "    with torch.no_grad():\n",
    "        model_output = model(**{k: v.to(model.device) for k, v in t.items()})\n",
    "    embeddings = model_output.last_hidden_state[:, 0, :]\n",
    "    embeddings = torch.nn.functional.normalize(embeddings)\n",
    "    return embeddings[0].cpu().squeeze()\n",
    "\n",
    "# Загрузка предварительно вычисленных векторов\n",
    "embeddings = np.loadtxt('embeddings.txt')\n",
    "embeddings_tensor = [torch.tensor(embedding) for embedding in embeddings]\n",
    "\n",
    "# Создание индекса Faiss\n",
    "embeddings_matrix = np.stack(embeddings)\n",
    "index = faiss.IndexFlatIP(embeddings_matrix.shape[1])\n",
    "index.add(embeddings_matrix)\n",
    "\n",
    "# Текст запроса\n",
    "text = 'добрую сказку с плохим концом для детей'\n",
    "\n",
    "# Встраивание запроса и поиск ближайших векторов с использованием Faiss\n",
    "query_embedding = embed_bert_cls(text)\n",
    "query_embedding = query_embedding.numpy().astype('float32')\n",
    "k, indices = index.search(np.expand_dims(query_embedding, axis=0), 3)\n",
    "\n",
    "# Вывод результатов\n",
    "for i in indices[0]:\n",
    "    print(df['title'][i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0aa7ef2-7f93-4300-9555-047bbc6c1036",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}