{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c307614e",
   "metadata": {},
   "source": [
    "### pretraining paraphrase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "eab76e07",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`torch_dtype` is deprecated! Use `dtype` instead!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initializing model from trained checkpoint: CLaRaConfig {\n",
      "  \"_attn_implementation_autoset\": true,\n",
      "  \"ae_mode\": \"token\",\n",
      "  \"attn_implementation\": null,\n",
      "  \"auto_map\": {\n",
      "    \"AutoConfig\": \"modeling_unirag.CLaRaConfig\",\n",
      "    \"AutoModel\": \"modeling_unirag.CLaRa\"\n",
      "  },\n",
      "  \"compr_base_model_name\": \"/mnt/ceph_rbd/model/Mistral-7B-Instruct-v0.2\",\n",
      "  \"compr_every_n_layer\": null,\n",
      "  \"compr_linear_type\": \"concat\",\n",
      "  \"compr_mlp_hidden_dim\": 8096,\n",
      "  \"compr_model_name\": null,\n",
      "  \"compr_n_layers\": 5,\n",
      "  \"compr_rate\": 16,\n",
      "  \"compr_rms_norm\": false,\n",
      "  \"compr_use_mlp\": false,\n",
      "  \"decoder_model_name\": \"/mnt/ceph_rbd/model/Mistral-7B-Instruct-v0.2\",\n",
      "  \"device_map\": null,\n",
      "  \"different_mem_tokens\": true,\n",
      "  \"doc_max_length\": 256,\n",
      "  \"generation_top_k\": 1,\n",
      "  \"kbtc_training\": false,\n",
      "  \"load_adapters\": false,\n",
      "  \"load_pretrained_checkpoint\": false,\n",
      "  \"lora\": true,\n",
      "  \"lora_compressor\": false,\n",
      "  \"lora_r\": 16,\n",
      "  \"lora_r_compressor\": 16,\n",
      "  \"max_new_tokens\": 128,\n",
      "  \"model_type\": \"CLaRa\",\n",
      "  \"optimize_mem_tokens\": true,\n",
      "  \"pad_token_id\": 2,\n",
      "  \"pure_inference\": false,\n",
      "  \"quantization\": \"no\",\n",
      "  \"sep\": true,\n",
      "  \"stage2_retrieval_top_n\": 1,\n",
      "  \"training_form\": \"both_separately\",\n",
      "  \"training_stage\": \"stage1\",\n",
      "  \"transformers_version\": \"4.56.2\"\n",
      "}\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading checkpoint shards: 100%|██████████| 3/3 [00:16<00:00,  5.57s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Base decoder parameters: 7241732096\n",
      "Model adapter keys: []\n",
      "Memory token count: 16\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The new embeddings will be initialized from a multivariate normal distribution that has old embeddings' mean and covariance. As described in this article: https://nlp.stanford.edu/~johnhew/vocab-expansion.html. To disable this, use `mean_resizing=False`\n",
      "The new lm_head weights will be initialized from a multivariate normal distribution that has old embeddings' mean and covariance. As described in this article: https://nlp.stanford.edu/~johnhew/vocab-expansion.html. To disable this, use `mean_resizing=False`\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading checkpoint adapter: decoder_adapter\n",
      "Loading checkpoint adapter: encoder_adapter\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/cprag/lib/python3.10/site-packages/peft/tuners/tuners_utils.py:196: UserWarning: Already found a `peft_config` attribute in the model. This will lead to having multiple adapters in the model. Make sure to know what you are doing!\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated answer ['The genus Weldenia belongs to the orchid family, Orchidaceae, and is native to regions in Mexico and Guatemala. It is a monotypic genus, meaning it contains only one species, Weldenia candida. This species is characterized by its succulent, stolon']\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoModel\n",
    "\n",
    "unirag = AutoModel.from_pretrained('/mnt/ceph_rbd/model/CLaRa-7B-Base/compression-16',trust_remote_code=True).to('cuda')\n",
    "\n",
    "# Example documents and question:\n",
    "documents = [\n",
    "    [\n",
    "        \"Weldenia is a monotypic genus of flowering plant in the family Commelinaceae, first describ ed in 1829. It has one single species: Weldenia candida, which grows originally in Mexico and Guatemala.\",\n",
    "        \"Hagsatera is a genus of flowering plants from the orchid family, Orchidaceae. There are two known species, native to Mexico and Guatemala\",\n",
    "        \"Alsobia is a genus of flowering plants in the family Gesneriaceae, native to Mexico, Guatemala and Costa Rica. The two species are succulent, stoloniferous herbs and were previously included in the genus \\\"Episcia\\\". Recent molecular studies have supported the separation of \\\"Alsobia\\\" from \\\"Episcia\\\"\"\n",
    "    ]\n",
    "]\n",
    "\n",
    "questions = [\"\" for _ in range(len(documents))]\n",
    "\n",
    "# End-to-end usage\n",
    "out = unirag.generate_from_paraphrase(questions=questions, documents=documents, max_new_tokens=64)\n",
    "print('Generated answer', out)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbcb5b73",
   "metadata": {},
   "source": [
    "### instruction tuning inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f89948d7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initializing model from trained checkpoint: CLaRaConfig {\n",
      "  \"_attn_implementation_autoset\": true,\n",
      "  \"ae_mode\": \"token\",\n",
      "  \"attn_implementation\": null,\n",
      "  \"auto_map\": {\n",
      "    \"AutoConfig\": \"modeling_unirag.CLaRaConfig\",\n",
      "    \"AutoModel\": \"modeling_unirag.CLaRa\"\n",
      "  },\n",
      "  \"compr_base_model_name\": \"/mnt/ceph_rbd/model/Mistral-7B-Instruct-v0.2\",\n",
      "  \"compr_every_n_layer\": null,\n",
      "  \"compr_linear_type\": \"concat\",\n",
      "  \"compr_mlp_hidden_dim\": 8096,\n",
      "  \"compr_model_name\": null,\n",
      "  \"compr_n_layers\": 5,\n",
      "  \"compr_rate\": 16,\n",
      "  \"compr_rms_norm\": false,\n",
      "  \"compr_use_mlp\": false,\n",
      "  \"decoder_model_name\": \"/mnt/ceph_rbd/model/Mistral-7B-Instruct-v0.2\",\n",
      "  \"device_map\": null,\n",
      "  \"different_mem_tokens\": true,\n",
      "  \"doc_max_length\": 256,\n",
      "  \"generation_top_k\": 5,\n",
      "  \"kbtc_training\": false,\n",
      "  \"load_adapters\": false,\n",
      "  \"load_pretrained_checkpoint\": false,\n",
      "  \"lora\": true,\n",
      "  \"lora_compressor\": false,\n",
      "  \"lora_r\": 16,\n",
      "  \"lora_r_compressor\": 16,\n",
      "  \"max_new_tokens\": 128,\n",
      "  \"model_type\": \"CLaRa\",\n",
      "  \"optimize_mem_tokens\": true,\n",
      "  \"pad_token_id\": 2,\n",
      "  \"pure_inference\": false,\n",
      "  \"quantization\": \"no\",\n",
      "  \"sep\": true,\n",
      "  \"stage2_retrieval_top_n\": 1,\n",
      "  \"training_form\": \"both_separately\",\n",
      "  \"training_stage\": \"stage1_2\",\n",
      "  \"transformers_version\": \"4.56.2\"\n",
      "}\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading checkpoint shards: 100%|██████████| 3/3 [00:00<00:00, 35.37it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Base decoder parameters: 7241732096\n",
      "Loading decoder adapter for stage1_2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model adapter keys: ['decoder_adapter']\n",
      "Memory token count: 16\n",
      "Adapter decoder_adapter trainable parameters: 41943040\n",
      "Loading checkpoint adapter: encoder_adapter\n",
      "Generated answer ['The Alsobia genus of plants grows originally in Mexico and Guatemala. However, based on the provided background information, the correct answer is Alsobia or Weldenia, the answer is Alsobia.']\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoModel\n",
    "\n",
    "unirag = AutoModel.from_pretrained('/mnt/ceph_rbd/model/CLaRa-7B-Instruct/compression-16',trust_remote_code=True).to('cuda')\n",
    "\n",
    "# Example documents and question:\n",
    "documents = [\n",
    "    [\n",
    "        \"Weldenia is a monotypic genus of flowering plant in the family Commelinaceae, first describ ed in 1829. It has one single species: Weldenia candida, which grows originally in Mexico and Guatemala.\",\n",
    "        \"Hagsatera is a genus of flowering plants from the orchid family, Orchidaceae. There are two known species, native to Mexico and Guatemala\",\n",
    "        \"Alsobia is a genus of flowering plants in the family Gesneriaceae, native to Mexico, Guatemala and Costa Rica. The two species are succulent, stoloniferous herbs and were previously included in the genus \\\"Episcia\\\". Recent molecular studies have supported the separation of \\\"Alsobia\\\" from \\\"Episcia\\\"\"\n",
    "    ]\n",
    "]\n",
    "\n",
    "questions = [\"Which genus of plant grows originally in Mexico and Guatemala, Phylica or Weldenia?\"]\n",
    "\n",
    "# instruction tuning inference \n",
    "out = unirag.generate_from_text(questions=questions, documents=documents, max_new_tokens=64)\n",
    "print('Generated answer', out)\n",
    "\n",
    "# "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be26f042",
   "metadata": {},
   "source": [
    "### end-to-end inference "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "efdf2b62",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initializing model from trained checkpoint: CLaRaConfig {\n",
      "  \"_attn_implementation_autoset\": true,\n",
      "  \"ae_mode\": \"token\",\n",
      "  \"attn_implementation\": null,\n",
      "  \"auto_map\": {\n",
      "    \"AutoConfig\": \"modeling_unirag.CLaRaConfig\",\n",
      "    \"AutoModel\": \"modeling_unirag.CLaRa\"\n",
      "  },\n",
      "  \"compr_base_model_name\": \"/mnt/ceph_rbd/model/Mistral-7B-Instruct-v0.2\",\n",
      "  \"compr_every_n_layer\": null,\n",
      "  \"compr_linear_type\": \"concat\",\n",
      "  \"compr_mlp_hidden_dim\": 8096,\n",
      "  \"compr_model_name\": null,\n",
      "  \"compr_n_layers\": 5,\n",
      "  \"compr_rate\": 16,\n",
      "  \"compr_rms_norm\": false,\n",
      "  \"compr_use_mlp\": false,\n",
      "  \"decoder_model_name\": \"/mnt/ceph_rbd/model/Mistral-7B-Instruct-v0.2\",\n",
      "  \"device_map\": null,\n",
      "  \"different_mem_tokens\": true,\n",
      "  \"doc_max_length\": 256,\n",
      "  \"generation_top_k\": 5,\n",
      "  \"kbtc_training\": false,\n",
      "  \"load_adapters\": false,\n",
      "  \"load_pretrained_checkpoint\": false,\n",
      "  \"lora\": true,\n",
      "  \"lora_compressor\": false,\n",
      "  \"lora_r\": 16,\n",
      "  \"lora_r_compressor\": 16,\n",
      "  \"max_new_tokens\": 128,\n",
      "  \"model_type\": \"CLaRa\",\n",
      "  \"optimize_mem_tokens\": true,\n",
      "  \"pad_token_id\": 2,\n",
      "  \"pure_inference\": false,\n",
      "  \"quantization\": \"no\",\n",
      "  \"sep\": true,\n",
      "  \"stage2_retrieval_top_n\": 1,\n",
      "  \"training_form\": \"both_separately\",\n",
      "  \"training_stage\": \"stage2\",\n",
      "  \"transformers_version\": \"4.56.2\"\n",
      "}\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading checkpoint shards: 100%|██████████| 3/3 [00:00<00:00, 35.89it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Base decoder parameters: 7241732096\n",
      "Model adapter keys: []\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory token count: 16\n",
      "Loading checkpoint adapter: decoder_adapter\n",
      "Loading checkpoint adapter: encoder_adapter\n",
      "Loading checkpoint adapter: query_reasoner_adapter\n",
      "Generated answer (['Weldenia'], tensor([[1, 0, 2, 4, 3]], device='cuda:0'))\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoModel\n",
    "\n",
    "unirag = AutoModel.from_pretrained('/mnt/ceph_rbd/model/CLaRa-7B-E2E/compression-16',trust_remote_code=True).to('cuda')\n",
    "\n",
    "# Example documents and question:\n",
    "# documents = [\n",
    "#     [x\n",
    "#         \"Weldenia is a monotypic genus of flowering plant in the family Commelinaceae, first describ ed in 1829. It has one single species: Weldenia candida, which grows originally in Mexico and Guatemala.\",\n",
    "#         \"Hagsatera is a genus of flowering plants from the orchid family, Orchidaceae. There are two known species, native to Mexico and Guatemala\",\n",
    "#         \"Alsobia is a genus of flowering plants in the family Gesneriaceae, native to Mexico, Guatemala and Costa Rica. The two species are succulent, stoloniferous herbs and were previously included in the genus \\\"Episcia\\\". Recent molecular studies have supported the separation of \\\"Alsobia\\\" from \\\"Episcia\\\"\"\n",
    "#     ]\n",
    "# ]\n",
    "\n",
    "documents = [[ \"Weldenia is a monotypic genus of flowering plant in the family Commelinaceae, first describ ed in 1829. It has one single species: Weldenia candida, which grows originally in Mexico and Guatemala.\" for _ in range(20)]]\n",
    "questions = [\"Which genus of plant grows originally in Mexico and Guatemala, Phylica or Weldenia?\"]\n",
    "\n",
    "# instruction tuning inference \n",
    "#the k is decided by the generation_top_k in the config.json\n",
    "out = unirag.generate_from_questions(questions=questions, documents=documents, max_new_tokens=64)\n",
    "print('Generated answer', out)\n",
    "\n",
    "# "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cprag",
   "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.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}