Upload faiss.ipynb

faiss.ipynb

@@ -0,0 +1,102 @@
+{
+ "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
+}