Upload 43 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+db/chroma.sqlite3 filter=lfs diff=lfs merge=lfs -text

app.ipynb

@@ -0,0 +1,274 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "276ff8dc-703d-4966-918b-983c592e7938",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "__import__('pysqlite3')\n",
+    "import sys\n",
+    "sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')\n",
+    "import chromadb\n",
+    "import torch\n",
+    "import gradio as gr\n",
+    "from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification, pipeline\n",
+    "\n",
+    "from langchain.llms import OpenAI, GigaChat\n",
+    "from langchain.chains import LLMChain\n",
+    "from langchain.prompts import PromptTemplate\n",
+    "\n",
+    "chatgpt = OpenAI(\n",
+    "    api_key='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',\n",
+    ")\n",
+    "\n",
+    "gigachat = GigaChat(\n",
+    "    credentials='Y2Y4Yjk5ODUtNThmMC00ODdjLTk5ODItNDdmYzhmNDdmNzE0OjQ5Y2RjNTVkLWFmMGQtNGJlYy04OGNiLTI1Yzc3MmJkMzYwYw==',\n",
+    "    scope='GIGACHAT_API_PERS',\n",
+    "    verify_ssl_certs=False\n",
+    ")\n",
+    "\n",
+    "llms = {\n",
+    "    'ChatGPT': chatgpt,\n",
+    "    'GigaChat': gigachat,\n",
+    "}\n",
+    "\n",
+    "# задаем формат вывода модели\n",
+    "answer_task_types = {\n",
+    "    'Развернутый ответ': 'Ответь достаточно подробно, но не используй ничего лишнего.',\n",
+    "    'Только цифры штрафа': 'Ответь в виде <количество> рублей или <диапазон> рублей, и больше ничего не пиши.'\n",
+    "}\n",
+    "\n",
+    "# проверяем с помощью LLM валидность запроса, исключая обработку бессмысленного входа\n",
+    "validity_template = '{query}\\n\\nЭто валидный запрос? Ответь да или нет, больше ничего не пиши.'\n",
+    "validity_prompt = PromptTemplate(template=validity_template, input_variables=['query'])\n",
+    "\n",
+    "# получаем ответ модели на запрос, используем его для более качественного поиска Retriever'ом и Cross-Encoder'ом\n",
+    "query_template = '{query} Ответь текстом, похожим на закон, не пиши ничего лишнего. Не используй в ответе слово КоАП РФ. Не используй слово \"Россия\".'\n",
+    "query_prompt = PromptTemplate(template=query_template, input_variables=['query'])\n",
+    "\n",
+    "# просим LLM выбрать один из 3 фрагментов текста, выбранных поисковыми моделями, где по мнению модели есть ответ. Если ответа нет, модель нам об этом сообщает\n",
+    "choose_answer_template = '1. {text_1}\\n\\n2. {text_2}\\n\\n3. {text_3}\\n\\nЗадание: выбери из перечисленных выше отрывков тот, где есть ответ на вопрос: \"{query}\". В качестве ответа напиши только номер 1, 2 или 3 и все. Если в данных отрывках нет ответа, то напиши \"Нет ответа\".'\n",
+    "choose_answer_prompt = PromptTemplate(template=choose_answer_template, input_variables=['text_1', 'text_2', 'text_3', 'query'])\n",
+    "\n",
+    "# просим LLM ответить на вопрос, опираясь на найденный фрагмент, и в нужном формате, или сообщить, что ответа все-таки нет\n",
+    "answer_template = '{text}\\n\\nЗадание: ответь на вопрос по тексту: \"{query}\". {answer_type} Если в данном тексте нет ответа, то напиши \"Нет ответа\".'\n",
+    "answer_prompt = PromptTemplate(template=answer_template, input_variables=['text', 'query', 'answer_type'])\n",
+    "\n",
+    "client = chromadb.PersistentClient(path='db')\n",
+    "collection = client.get_collection(name=\"administrative_codex\")\n",
+    "\n",
+    "retriever_checkpoint = 'sentence-transformers/LaBSE'\n",
+    "retriever_tokenizer = AutoTokenizer.from_pretrained(retriever_checkpoint)\n",
+    "retriever_model = AutoModel.from_pretrained(retriever_checkpoint)\n",
+    "\n",
+    "cross_encoder_checkpoint = 'jeffwan/mmarco-mMiniLMv2-L12-H384-v1'\n",
+    "cross_encoder_model = AutoModelForSequenceClassification.from_pretrained(cross_encoder_checkpoint)\n",
+    "cross_encoder_tokenizer = AutoTokenizer.from_pretrained(cross_encoder_checkpoint)\n",
+    "cross_encoder = pipeline('text-classification', model=cross_encoder_model, tokenizer=cross_encoder_tokenizer)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "ca60b616-97f4-4d01-b0f5-4a994d11e216",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Running on local URL:  http://127.0.0.1:7864\n",
+      "\n",
+      "To create a public link, set `share=True` in `launch()`.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div><iframe src=\"http://127.0.0.1:7864/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": []
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "def encode(docs):\n",
+    "    if type(docs) == str:\n",
+    "        docs = [docs]\n",
+    "\n",
+    "    encoded_input = retriever_tokenizer(\n",
+    "        docs,\n",
+    "        padding=True,\n",
+    "        truncation=True,\n",
+    "        max_length=512,\n",
+    "        return_tensors='pt'\n",
+    "    )\n",
+    "    \n",
+    "    with torch.no_grad():\n",
+    "        model_output = retriever_model(**encoded_input)\n",
+    "        \n",
+    "    embeddings = model_output.pooler_output\n",
+    "    embeddings = torch.nn.functional.normalize(embeddings)\n",
+    "    return embeddings.detach().cpu().tolist()\n",
+    "\n",
+    "\n",
+    "def re_rank(sentence, docs):\n",
+    "    return [res['score'] for res in cross_encoder([{'text': sentence, 'text_pair': doc} for doc in docs], max_length=512, truncation=True)]\n",
+    "\n",
+    "\n",
+    "def update_query_with_llm(query, llm_type, use_llm_for_retriever):\n",
+    "    if llm_type == 'Без LLM' or not use_llm_for_retriever:\n",
+    "        return query\n",
+    "    \n",
+    "    llm_chain = LLMChain(prompt=query_prompt, llm=llms[llm_type])\n",
+    "    return f'{query} {llm_chain.run(query).strip()}'\n",
+    "\n",
+    "\n",
+    "def answer_with_llm(query, re_ranked_res, llm_type, llm_answer_type):\n",
+    "    if llm_type == 'Без LLM':\n",
+    "        answer, metadata, re_ranker_score = re_ranked_res[0]\n",
+    "    else:\n",
+    "        llm_chain = LLMChain(prompt=choose_answer_prompt, llm=llms[llm_type])\n",
+    "        llm_chain_dict = {f'text_{i}': res[0] for i, res in enumerate(re_ranked_res, start=1)}\n",
+    "        llm_chain_dict['query'] = query\n",
+    "\n",
+    "        llm_res = llm_chain.run(llm_chain_dict).strip()\n",
+    "\n",
+    "        if 'нет ответа' in llm_res.lower() or not llm_res[0].isnumeric():\n",
+    "            return 'Нет ответа', '', ''\n",
+    "        \n",
+    "        most_suitable_text, metadata, re_ranker_score = re_ranked_res[int(llm_res[0]) - 1]\n",
+    "\n",
+    "        llm_chain = LLMChain(prompt=answer_prompt, llm=llms[llm_type])\n",
+    "        answer = llm_chain.run({'text': most_suitable_text, 'query': query, 'answer_type': llm_answer_type}).strip()\n",
+    "\n",
+    "        if 'нет ответа' in answer.lower():\n",
+    "            answer = 'Нет ответа'\n",
+    "\n",
+    "    # если LLM сначала выбрала фрагмент, где есть ответ, а потом не смогла ответить на вопрос (что бывает редко), то все равно порекомендуем пользователю обратиться к норме\n",
+    "    law_norm = f\"{'Попробуйте обратиться к этому источнику: ' if answer == 'Нет ответа' else ''}{metadata['article']} {metadata['point']} {metadata['doc']}\"\n",
+    "    return answer, law_norm, re_ranker_score\n",
+    "\n",
+    "\n",
+    "def check_request_validity(func):\n",
+    "    def wrapper(\n",
+    "        query,\n",
+    "        llm_type,\n",
+    "        llm_answer_type,\n",
+    "        use_llm_for_retriever,\n",
+    "        use_llm_for_request_validation\n",
+    "    ):\n",
+    "        query = query.strip()\n",
+    "        \n",
+    "        if not query:\n",
+    "            return 'Невалидный запрос', '', ''\n",
+    "            \n",
+    "        if llm_type == 'Без LLM' or not use_llm_for_request_validation:\n",
+    "            return func(query, llm_type, llm_answer_type, use_llm_for_retriever)\n",
+    "        \n",
+    "        llm_chain = LLMChain(prompt=validity_prompt, llm=llms[llm_type])\n",
+    "\n",
+    "        if 'нет' in llm_chain.run(query).lower():\n",
+    "            return 'Невалидный запрос', '', ''\n",
+    "            \n",
+    "        return func(query, llm_type, llm_answer_type, use_llm_for_retriever)\n",
+    "    \n",
+    "    return wrapper\n",
+    "\n",
+    "\n",
+    "@check_request_validity\n",
+    "def fn(\n",
+    "    query,\n",
+    "    llm_type,\n",
+    "    llm_answer_type,\n",
+    "    use_llm_for_retriever\n",
+    "):\n",
+    "    # обогатим запрос с помощью LLM, чтобы поис��овым моделям было проще найти нужный фрагмент с ответом\n",
+    "    retriever_ranker_query = update_query_with_llm(query, llm_type, use_llm_for_retriever)\n",
+    "\n",
+    "    # Retriever-поиск по базе данных\n",
+    "    retriever_res = collection.query(\n",
+    "        query_embeddings=encode(retriever_ranker_query),\n",
+    "        n_results=10,\n",
+    "    )\n",
+    "\n",
+    "    top_k_docs = retriever_res['documents'][0]\n",
+    "\n",
+    "    # re-ranking с помощью Cross-Encoder'а и отбор лучших кандидатов\n",
+    "    re_rank_scores = re_rank(retriever_ranker_query, top_k_docs)\n",
+    "    re_ranked_res = sorted(\n",
+    "        [[doc, meta, score] for doc, meta, score in zip(retriever_res['documents'][0], retriever_res['metadatas'][0], re_rank_scores)],\n",
+    "        key=lambda x: x[-1],\n",
+    "        reverse=True,\n",
+    "    )[:3]\n",
+    "\n",
+    "    # поиск ответа и нормы с помощью LLM\n",
+    "    return answer_with_llm(query, re_ranked_res, llm_type, llm_answer_type)\n",
+    "\n",
+    "\n",
+    "demo = gr.Interface(\n",
+    "    fn=fn,\n",
+    "    inputs=[\n",
+    "        gr.Textbox(lines=3, label='Запрос', placeholder='Введите запрос'),\n",
+    "        gr.Dropdown(label='Тип LLM', choices=['ChatGPT', 'GigaChat', 'Без LLM'], value='ChatGPT'),\n",
+    "        gr.Dropdown(label='Тип итогового ответа LLM', choices=['Только цифры штрафа', 'Развернутый ответ'], value='Только цифры штрафа'),\n",
+    "        gr.Checkbox(label=\"Использовать LLM для Retriever'а\", value=True, info=\"При использовании LLM для Retriever'а ко входному запросу будет добавляться промежуточный ответ LLM на запрос. Это способствует повышению качества поиска ответа.\"),\n",
+    "        gr.Checkbox(label=\"Использовать LLM для проверки валидности запроса\", value=False)\n",
+    "    ],\n",
+    "    outputs=[\n",
+    "        gr.Textbox(label='Ответ'),\n",
+    "        gr.Textbox(label='Норма'),\n",
+    "        gr.Textbox(label=\"Уверенность Cross-Encoder'а\"),\n",
+    "    ],\n",
+    ")\n",
+    "\n",
+    "demo.launch()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "17ebff84-9312-46ae-9972-cd859fad36f9",
+   "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.8.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

db/97372335-8399-4edb-90ed-bcaf83a48591/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e861bdd7817a442cd756c01729685476c6d03bf7320edfa45a9d31217dc60f89
+size 6424000

db/97372335-8399-4edb-90ed-bcaf83a48591/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:995285888fae05329c795e15bef4387c9003f7ed65af6bd46fc6a1220b9e2265
+size 100

db/97372335-8399-4edb-90ed-bcaf83a48591/index_metadata.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:659f72af1aa9a4dc9bba946c9c62147869d7b1eaef779684a67821dd6ba565bf
+size 52870

db/97372335-8399-4edb-90ed-bcaf83a48591/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:135cb2913a715740c00acafbd2a392317f77955f782a1e2974ded7c7f0d769e6
+size 8000

db/97372335-8399-4edb-90ed-bcaf83a48591/link_lists.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4c44a130f362f4e7bd70389814b9d8d8889d0d1262ca68d6c5fd7ae2fb11e707
+size 16976

db/9dc5e8eb-9345-462f-a1fa-938aec17a9a0/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aef8dfed4c15fb14fd74ce9f26f2aecf739df73b7ae473f2e4a36baa830770d0
+size 6424000

db/9dc5e8eb-9345-462f-a1fa-938aec17a9a0/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:91a34e115568431e60861fa63dee2fbc5eba7750376a27655abbd690383acc8a
+size 100

db/9dc5e8eb-9345-462f-a1fa-938aec17a9a0/index_metadata.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e015a6c87c458effc1fb3105fbdb92b3681df738b60f2d64624af959b206afe0
+size 52870

db/9dc5e8eb-9345-462f-a1fa-938aec17a9a0/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dd2b7f4046e646fb88c6d837cb3d2067ba3b0d82f256c7f3afe4a5fd4f26a655
+size 8000

db/9dc5e8eb-9345-462f-a1fa-938aec17a9a0/link_lists.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6209ce61c94144576d5851341ac04d85cf6b558c02da410e983b3f9290e411bc
+size 16976

db/aa5bd886-4fdf-4d66-965d-18ffddafee21/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d1e1753e5819ae5dffa99158015f31ba066c12ee0e4f657e64dc64660dd7b77f
+size 6424000

db/aa5bd886-4fdf-4d66-965d-18ffddafee21/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:995285888fae05329c795e15bef4387c9003f7ed65af6bd46fc6a1220b9e2265
+size 100

db/aa5bd886-4fdf-4d66-965d-18ffddafee21/index_metadata.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0086f16ab6934ed3c2a5b99a1e9595e150623259c8396e0faf2f26023f20c757
+size 52870

db/aa5bd886-4fdf-4d66-965d-18ffddafee21/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:38c55df27ec0a8adb0ac1e2d55377a84ddbb939c3f3f9704a150952b067108e5
+size 8000

db/aa5bd886-4fdf-4d66-965d-18ffddafee21/link_lists.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6de5ffb84a9d65860195295c241ed10a8618676c290dfa8379158bef5c308f5e
+size 16976

db/b402ad0c-87fe-4eff-b65e-58777a7face5/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b72c643cdc6cdf14f5bf99407c6df138da0ce720c6b2127ee36aea4fe57a3b50
+size 6424000

db/b402ad0c-87fe-4eff-b65e-58777a7face5/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:995285888fae05329c795e15bef4387c9003f7ed65af6bd46fc6a1220b9e2265
+size 100

db/b402ad0c-87fe-4eff-b65e-58777a7face5/index_metadata.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7ddf8414b24f52bb480445ed59c48b82cd02324d2668f74f8582e8d16d1558ab
+size 52870

db/b402ad0c-87fe-4eff-b65e-58777a7face5/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:287895f1d9f79674f38d9748c5e8a38edef18d74567d51b7790eb3094d959605
+size 8000

db/b402ad0c-87fe-4eff-b65e-58777a7face5/link_lists.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0c4b1a82fceed2e2b272bffdc76f6dd30062dfa5fcdf3ca1292fd739a764cfdb
+size 16976

db/chroma.sqlite3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c29da25687cacaa124ee55421db14f74f971e5e185a029e7113f3d047bac697f
+size 29159424

db/e90dea15-53e7-4b41-b502-0ab854ae08d8/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:02d382a284aa159c52564b0fc03080a00b4ecbbbbe383781d81f004a21c7e61a
+size 6424000

db/e90dea15-53e7-4b41-b502-0ab854ae08d8/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:995285888fae05329c795e15bef4387c9003f7ed65af6bd46fc6a1220b9e2265
+size 100

db/e90dea15-53e7-4b41-b502-0ab854ae08d8/index_metadata.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ccf1f97bd21b051c83e0e8a19d1e4f7de0972d3095b7b1807f94dfeafc99db89
+size 52870

db/e90dea15-53e7-4b41-b502-0ab854ae08d8/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cc538b666ab7415eda93fabfb42aa70ec82d57f19509da8eb899e10faa21a153
+size 8000

db/e90dea15-53e7-4b41-b502-0ab854ae08d8/link_lists.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:94b06e796bcc2b6f1959f2575e260cc9ed17950526af440de41f1bb7f148ed0f
+size 16976

preprocess_doc.ipynb

@@ -0,0 +1,220 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "bc1b8947-4397-4fb8-b3ae-310fdb44c056",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "__import__('pysqlite3')\n",
+    "import sys\n",
+    "import os\n",
+    "sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')\n",
+    "os.environ['ALLOW_RESET'] = 'True'\n",
+    "\n",
+    "import torch\n",
+    "from torch.utils.data import DataLoader\n",
+    "from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification, pipeline\n",
+    "import numpy as np\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "import chromadb"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c0af096a-f302-4df9-9f26-63213ad44b8f",
+   "metadata": {},
+   "source": [
+    "### Подготавливаем базу данных"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "e5a2886c-586a-49f7-8833-0972407bf1fa",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "client = chromadb.PersistentClient(path='db')\n",
+    "client.reset()\n",
+    "\n",
+    "collection = client.create_collection(\n",
+    "    name=\"administrative_codex\",\n",
+    "    metadata={\"hnsw:space\": \"cosine\"}\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "329c4f86-c514-4039-9494-621cb7042f77",
+   "metadata": {},
+   "source": [
+    "### Открываем и предобрабатываем КоАП"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "f057b53d-5a50-4a92-bbda-4c8396aca107",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open('docs/КоАП РФ.txt', encoding='utf-8') as r:\n",
+    "    raw_text = r.read().split('\\n\\n')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5589508c-5be9-403d-ba68-7a0e60658136",
+   "metadata": {},
+   "source": [
+    "### Делим документ по частям статей, исключаем лишнее"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "d7b87b89-5cc2-4c3b-8cee-8ea59b4778a5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "paragraphs = []\n",
+    "index = 0\n",
+    "\n",
+    "while index != len(raw_text):\n",
+    "    if raw_text[index].startswith('Статья'):\n",
+    "        article = ' '.join(raw_text[index].strip().split()[:2])\n",
+    "        article_points = raw_text[index + 1].split('\\n')\n",
+    "\n",
+    "        cur_point = ''\n",
+    "        for i in range(len(article_points)):\n",
+    "            cur_point_part = article_points[i].strip()\n",
+    "            \n",
+    "            if 'КонсультантПлюс' in article_points[i] + article_points[i - 1]:\n",
+    "                continue\n",
+    "            elif cur_point_part.split()[0].strip().replace('.', '').isnumeric() or cur_point_part.startswith('Примечание. '):\n",
+    "                if cur_point:\n",
+    "                    if cur_point.startswith('Примечание. '):\n",
+    "                        paragraphs.append([cur_point, article, 'Примечание.'])\n",
+    "                    elif cur_point[0].isnumeric():\n",
+    "                        paragraphs.append([' '.join(cur_point.split()[1:]), article, f'Часть {cur_point.split()[0]}'])\n",
+    "                    else:\n",
+    "                        paragraphs.append([cur_point, article, ''])\n",
+    "                        \n",
+    "                cur_point = cur_point_part\n",
+    "            elif cur_point_part[0] != '(' and cur_point_part[-1] != ')' and 'утратил силу' not in cur_point_part[:20].lower():\n",
+    "                cur_point += ' ' + cur_point_part\n",
+    "            \n",
+    "        index += 2\n",
+    "    else:\n",
+    "        index += 1"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5e574c65-8100-45e3-9d8b-e7a15bfec242",
+   "metadata": {},
+   "source": [
+    "### Получаем эмбеддинги из извлеченных фрагментов и сохраняем их в базу данных"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "35204261-c2a0-43bc-b167-931154cfb77f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "checkpoint = 'sentence-transformers/LaBSE'\n",
+    "tokenizer = AutoTokenizer.from_pretrained(checkpoint)\n",
+    "model = AutoModel.from_pretrained(checkpoint, device_map='cuda:0')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "7aad08e2-fca0-4cf8-b568-17e5c2b4ffff",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def encode(docs):\n",
+    "    if type(docs) == str:\n",
+    "        docs = [docs]\n",
+    "\n",
+    "    encoded_input = tokenizer(\n",
+    "        docs,\n",
+    "        padding=True,\n",
+    "        truncation=True,\n",
+    "        max_length=512,\n",
+    "        return_tensors='pt'\n",
+    "    )\n",
+    "    \n",
+    "    with torch.no_grad():\n",
+    "        model_output = model(**encoded_input.to('cuda'))\n",
+    "        \n",
+    "    embeddings = model_output.pooler_output\n",
+    "    embeddings = torch.nn.functional.normalize(embeddings)\n",
+    "    return embeddings.detach().cpu().tolist()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "f118299e-c57a-40ce-a060-5a86b452cfec",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "BATCH_SIZE = 128\n",
+    "loader = DataLoader(paragraphs, batch_size=BATCH_SIZE)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "95dafd67-f6f8-41e1-a209-6325a23e066f",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 17/17 [00:11<00:00,  1.48it/s]\n"
+     ]
+    }
+   ],
+   "source": [
+    "for i, docs in enumerate(tqdm(loader)):\n",
+    "    embeddings = encode(docs[0])\n",
+    "    collection.add(\n",
+    "        documents=docs[0],\n",
+    "        metadatas=[{'doc': 'КоАП РФ', 'article': a, 'point': p} for a, p in zip(docs[1], docs[2])],\n",
+    "        embeddings=embeddings,\n",
+    "        ids=[f'id{i * BATCH_SIZE + j}' for j in range(len(docs[0]))],\n",
+    "    )"
+   ]
+  }
+ ],
+ "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.8.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

question_generation.ipynb

@@ -0,0 +1,179 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "d024645c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "__import__('pysqlite3')\n",
+    "import sys\n",
+    "import os\n",
+    "sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')\n",
+    "os.environ['ALLOW_RESET'] = 'True'\n",
+    "\n",
+    "import pandas as pd\n",
+    "from tqdm import tqdm\n",
+    "import time\n",
+    "\n",
+    "import chromadb\n",
+    "from gigachat import GigaChat\n",
+    "\n",
+    "client = chromadb.PersistentClient(path='db')\n",
+    "collection = client.get_collection(name=\"administrative_codex\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "17dae6a5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "docs = collection.get()['documents']\n",
+    "prompt = 'Задание: напиши в виде нумерованного списка 3 конкретных независимых друг от друга вопроса, ответ на которые можно найти в приведенном тексте. Не упоминай федеральные законы. Не упоминай КоАП.'"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "91549726-3c7a-44ef-8519-c1afc3adde0f",
+   "metadata": {},
+   "source": [
+    "### Генерируем вопросы к каждому фрагменту текста"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "id": "06f82948",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|███████████████████████████████████████████████████████████████████████████████| 2130/2130 [54:44<00:00,  1.54s/it]\n"
+     ]
+    }
+   ],
+   "source": [
+    "for doc in tqdm(docs[1:]):\n",
+    "    question_ready = False\n",
+    "    \n",
+    "    while not question_ready:\n",
+    "        try:\n",
+    "            text = f'{doc}\\n\\n{prompt}'\n",
+    "    \n",
+    "            with GigaChat(credentials='N2ZiNDIxZTgtM2Y4Yy00MGJjLWI4OTgtN2M5NGM5MTYzZTNiOmFmYjJmZTUwLTc1OWItNGQ5MC1iMGVmLTMwYTNlODU3YzVmZg==', scope='GIGACHAT_API_PERS', verify_ssl_certs=False) as giga:\n",
+    "                questions = giga.chat(text).choices[0].message.content\n",
+    "            \n",
+    "            question_ready = True\n",
+    "        except:\n",
+    "            time.sleep(5)\n",
+    "\n",
+    "    df = pd.read_csv('generated_questions.csv')\n",
+    "    new_df = pd.DataFrame({'text': [doc], 'questions': [questions]})\n",
+    "    pd.concat([df, new_df], ignore_index=True).to_csv('generated_questions.csv', index=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "id": "2f44eac2-7ce0-4d26-9f4a-41f5bfe0fa44",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "generated_questions_df = pd.read_csv('generated_questions.csv')\n",
+    "docs = generated_questions_df['text'].tolist()\n",
+    "generated_questions = generated_questions_df['questions'].tolist()\n",
+    "\n",
+    "prompt = 'В России. Дай подробный ответ текстом, похожим на закон, не пиши ничего лишнего.'"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "90b543f8-0b94-4c0c-8a69-9574b7c54db9",
+   "metadata": {},
+   "source": [
+    "### Генерируем ответы к вопросам, в которых есть слово штраф"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 39,
+   "id": "a7937078-6200-44ba-b43f-4867e947b750",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|███████████████████████████████████████████████████████████████████████████████| 1978/1978 [50:30<00:00,  1.53s/it]\n"
+     ]
+    }
+   ],
+   "source": [
+    "for doc, g_questions in zip(tqdm(docs[153:]), generated_questions[153:]):\n",
+    "    llm_answer_ready = False\n",
+    "    fine_question = ''\n",
+    "\n",
+    "    for question in g_questions.split('\\n'):\n",
+    "        question = ' '.join(question.split()[1:])\n",
+    "        \n",
+    "        if 'штраф' in question:\n",
+    "            fine_question = question\n",
+    "            break\n",
+    "\n",
+    "    if not fine_question:\n",
+    "        continue\n",
+    "    \n",
+    "    while not llm_answer_ready:\n",
+    "        try:\n",
+    "            text = f'Помоги, пожалуйста. {fine_question} {prompt}'\n",
+    "    \n",
+    "            with GigaChat(credentials='MmU3OTdhNmItMTQzYy00NGQzLWEyYTctZjcxOWJmYThiMWE5OmE1ZDdhNDkxLWI5ZTEtNGFkZS04N2JjLTExZjE5MTYwNGQ5Yg==', scope='GIGACHAT_API_PERS', verify_ssl_certs=False) as giga:\n",
+    "                llm_answer = giga.chat(text).choices[0].message.content.split('\\n')[0]\n",
+    "            \n",
+    "            llm_answer_ready = True\n",
+    "        except:\n",
+    "            time.sleep(5)\n",
+    "\n",
+    "    \n",
+    "    if len(llm_answer) > 100:\n",
+    "        df = pd.read_csv('generated_additional_llm_answer.csv')\n",
+    "        new_df = pd.DataFrame({'text': [doc], 'question': [fine_question], 'llm_answer': [llm_answer]})\n",
+    "        pd.concat([df, new_df], ignore_index=True).to_csv('generated_additional_llm_answer.csv', index=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5fd64855-01b5-4c66-a425-b6d91b355a22",
+   "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.8.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}