app.ipynb

{
 "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
}