{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer ,AutoModelForCausalLM\n",
    "import torch\n",
    "import re\n",
    "from sklearn.utils import shuffle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [],
   "source": [
    "# util function to get expected len after tokenizing\n",
    "def get_length_param(text: str, tokenizer) -> str:\n",
    "    tokens_count = len(tokenizer.encode(text))\n",
    "    if tokens_count <= 15:\n",
    "        len_param = '1'\n",
    "    elif tokens_count <= 50:\n",
    "        len_param = '2'\n",
    "    elif tokens_count <= 256:\n",
    "        len_param = '3'\n",
    "    else:\n",
    "        len_param = '-'\n",
    "    return len_param"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "outputs": [],
   "source": [
    "def remove_duplicates(S):\n",
    "    S = re.sub(r'[a-zA-Z]+', '', S) #Remove english\n",
    "    S = S.split()\n",
    "    result = \"\"\n",
    "    for subst in S:\n",
    "        if subst not in result:\n",
    "            result += subst+\" \"\n",
    "    return result.rstrip()"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [],
   "source": [
    "fit_checkpoint = \"WarBot\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(fit_checkpoint)\n",
    "model = AutoModelForCausalLM.from_pretrained(fit_checkpoint)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "outputs": [],
   "source": [
    "quote = \"Однажды мы проснёмся и поймём, что бригада Нахаль наваляла десантникам по самые помидоры\""
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "outputs": [],
   "source": [
    "# encode the input, add the eos_token and return a tensor in Pytorch\n",
    "user_inpit_ids = tokenizer.encode(f\"|0|{get_length_param(quote, tokenizer)}|\" \\\n",
    "                                              + quote + tokenizer.eos_token, return_tensors=\"pt\")\n",
    "\n",
    "#chat_history_ids = torch.cat([chat_history_ids, user_inpit_ids], dim=-1)\n",
    "\n",
    "chat_history_ids = user_inpit_ids # To be changed\n",
    "\n",
    "output_id = model.generate(\n",
    "            chat_history_ids,\n",
    "            num_return_sequences=1, # use for more variants, but have to print [i]\n",
    "            max_length=300, #512\n",
    "            no_repeat_ngram_size=1, #3\n",
    "            do_sample=True, #True\n",
    "            top_k=50,#50\n",
    "            top_p=0.9, #0.9\n",
    "            temperature = 0.45, # was 0.6, 0 for greedy\n",
    "            #mask_token_id=tokenizer.mask_token_id,\n",
    "            eos_token_id=tokenizer.eos_token_id,\n",
    "            #unk_token_id=tokenizer.unk_token_id,\n",
    "            pad_token_id=tokenizer.pad_token_id,\n",
    "            #pad_token_id=tokenizer.eos_token_id,\n",
    "            #device='cpu'\n",
    "        )"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "outputs": [],
   "source": [
    "def removeSigns(S):\n",
    "    last_index = max(S.rfind(\".\"), S.rfind(\"!\"))\n",
    "    if last_index >= 0:\n",
    "        S = S[:last_index+1]\n",
    "    return S"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "outputs": [],
   "source": [
    "def getResponce():\n",
    "    response = tokenizer.decode(output_id[0], skip_special_tokens=True)\n",
    "    response = removeSigns(response)\n",
    "    #response = re.sub(r'[^а-яА-Я;.,!?]', '', response) # Clear the response, remains only russian\n",
    "    response_с = response.split(quote)[-1] #Remove the Quote\n",
    "    clean_response = remove_duplicates(re.sub(r\"\\d{4,}\", \"\", response_с)) # Remove the consequent numbers with 4 or more digits\n",
    "    return clean_response"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Response: в бригаде есть несколько батальонов \"йерихон\". они основном для того чтобы отражать атаки хезов. батальонный уровень это ротные минометы на джипах прицепами (на уровне батальона). если них будет достаточно ракет могут даже накрыть батарею 120мм минометной установки или из состава бригады может быть хуже чем батареи 122 м-109 которые находятся под управлением роты/бат аля рейнджеры... опять все зависит ситуации например после первой ливанской артиллеристы стали очень сильно нервничать когда обстреливали израильские бпла типа 28, как приходилось отвечать свои задачи. теперь вот примеру американцы решили полностью перевести всю бригаду второй эшелон : 1) сократив количество артдивизионов 4х; 3 пехотных батальонах(м113); 5 танковых + отдельный армейский который сможет прикрывать танки непосредственно перед атакой противника.. правда нужно еще иметь возможность поддерживать свой штатную авиацию огнем своего штатного места без необходимости перебрасывать туда часть танков.... короче говоря вся эта система должна работать вместе /при условии полного отсутствия взаимозачетчиков между ними...но тут надо смотреть кто первый окажется дежурным батареей техасских коптеров..и итп....\n"
     ]
    }
   ],
   "source": [
    "print(\"Response:\",getResponce())"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Spelling Fix:"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "outputs": [],
   "source": [
    "from autocorrect import Speller\n",
    "spell = Speller('ru')\n",
    "spell_fix_response = spell(getResponce())"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "в бригаде есть несколько батальонов \"рихон\". они основном для того чтобы отражать атаки уезов. батальонный уровень это ротные минометы на джипах прицепами (на уровне батальона). если них будет достаточно ракет могут даже накрыть батарею 120мм минометной установки или из состава бригады может быть хуже чем батареи 122 м-109 которые находятся под управлением роты/бат аля рейнджеры... опять все зависит ситуации например после первой ливанской артиллеристы стали очень сильно нервничать когда обстреливали израильские была типа 28, как приходилось отвечать свои задачи. теперь вот примеру американцы решили полностью перевести всю бригаду второй эшелон : 1) сократив количество артдивизионов 4х; 3 пехотных батальонах(м113); 5 танковых + отдельный армейский который сможет прикрывать танки непосредственно перед атакой противника.. правда нужно еще иметь возможность поддерживать свой штатную авиацию огнем своего штатного места без необходимости перебрасывать туда часть танков.... короче говоря вся эта система должна работать вместе /при условии полного отсутствия взаимозачетчиков между ними...но тут надо смотреть кто первый окажется дежурным батареей техасских коптеров..и итп....\n"
     ]
    }
   ],
   "source": [
    "print(spell_fix_response)"
   ],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}