doc/tree_summarize.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Summarize the paper as a tree\n",
    "In this example we read the labeled paper as a tree.\n",
    "Then we summarize it by using the LlamaIndex TreeSummarize function.\n",
    "\n",
    "We can either use OpenAI's API or load a model to do it locally."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load LLM and TreeSummarize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/duan/texcompile/venv/lib/python3.8/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "Loading checkpoint shards: 100%|██████████| 2/2 [00:06<00:00,  3.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "******\n",
      "Could not load OpenAIEmbedding. Using HuggingFaceBgeEmbeddings with model_name=BAAI/bge-small-en. If you intended to use OpenAI, please check your OPENAI_API_KEY.\n",
      "Original error:\n",
      "No API key found for OpenAI.\n",
      "Please set either the OPENAI_API_KEY environment variable or openai.api_key prior to initialization.\n",
      "API keys can be found or created at https://platform.openai.com/account/api-keys\n",
      "\n",
      "******\n"
     ]
    }
   ],
   "source": [
    "from llama_index.response_synthesizers import TreeSummarize\n",
    "from llama_index import SimpleDirectoryReader, ServiceContext\n",
    "import os\n",
    "os.chdir('..')\n",
    "#import openai\n",
    "#openai.api_key = \"your key\"\n",
    "\n",
    "import torch\n",
    "from llama_index.llms import HuggingFaceLLM\n",
    "from llama_index.prompts import PromptTemplate\n",
    "\n",
    "system_prompt = \"\"\"<|SYSTEM|># StableLM Tuned (Alpha version)\n",
    "- StableLM is a helpful and harmless open-source AI language model developed by StabilityAI.\n",
    "- StableLM is excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.\n",
    "- StableLM is more than just an information source, StableLM is also able to write poetry, short stories, and make jokes.\n",
    "- StableLM will refuse to participate in anything that could harm a human.\n",
    "\"\"\"\n",
    "\n",
    "# This will wrap the default prompts that are internal to llama-index\n",
    "query_wrapper_prompt = PromptTemplate(\"<|USER|>{query_str}<|ASSISTANT|>\")\n",
    "\n",
    "llm = HuggingFaceLLM(\n",
    "    context_window=4096,\n",
    "    max_new_tokens=256,\n",
    "    generate_kwargs={\"temperature\": 0.7, \"do_sample\": False},\n",
    "    system_prompt=system_prompt,\n",
    "    query_wrapper_prompt=query_wrapper_prompt,\n",
    "    tokenizer_name=\"StabilityAI/stablelm-tuned-alpha-3b\",\n",
    "    model_name=\"StabilityAI/stablelm-tuned-alpha-3b\",\n",
    "    device_map=\"auto\",\n",
    "    stopping_ids=[50278, 50279, 50277, 1, 0],\n",
    "    tokenizer_kwargs={\"max_length\": 4096},\n",
    "    # uncomment this if using CUDA to reduce memory usage\n",
    "    model_kwargs={\"torch_dtype\": torch.float16}\n",
    ")\n",
    "service_context = ServiceContext.from_defaults(chunk_size=1024, llm=llm)\n",
    "summarizer = TreeSummarize(service_context=service_context, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load data from exported csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "df = pd.read_csv('text.csv', sep='\\t')\n",
    "text = ''\n",
    "section_id = 0\n",
    "for i, row in df[df['reading_order']!=-1].sort_values(by=['reading_order', 'Unnamed: 0']).iterrows():\n",
    "    if row['section_id'] > section_id:\n",
    "        text += '\\n'\n",
    "        section_id = row['section_id']\n",
    "    if row['label'] != 'Figure':\n",
    "        text += row['token'] + ' '"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 text chunks after repacking\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 text chunks after repacking\n"
     ]
    }
   ],
   "source": [
    "sections =  text.split('\\n')\n",
    "response = summarizer.get_response(\"what are this texts talking about, emphasize the second section:\", sections)\n",
    "response2 = summarizer.get_response(\"what are this texts talking about, emphasize the evaluation section:\", sections)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'This text is discussing the use of Convolutional Neural Networks (CNNs) for automatic crater detection. Crater detection is a crucial step in the study of the Solar System, as it provides valuable information about the past and present geological processes. CNNs have been used successfully in the past to detect craters, but they have not yet been widely adopted for automatic crater detection. The authors present a method for automatic crater detection using CNNs, which are organized as a computation graph. The CNN architecture is a crucial component of the computation graph, as it allows for the efficient processing of large amounts of data. The authors also discuss the use of CNNs for detecting craters and the challenges they face in automatically detecting craters. The text highlights the importance of using CNNs for automatic crater detection and the potential benefits of using them for this purpose.'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'This text is discussing the Crater Detection via Convolutional Neural Networks (CNN) approach, which uses advanced machine learning techniques to automatically detect craters. The text highlights the challenges of automatically detecting craters, such as the complex surface of craters, and the need for advanced machine learning techniques to deal with the large amount of satellite imagery. The text also mentions the use of Convolutional Neural Networks (CNN) for automatic crater detection, which is a significant breakthrough in the field. The text also mentions the use of CNNs for automatic crater detection, which is a new approach that has achieved good performance on this dataset. The text also mentions the use of Banderia, a dataset that contains high-quality craters across all tiles, and the evaluation of the CNN approach using the Banderia dataset. The text concludes by discussing the potential of CNNs for automatic crater detection and the need for further advancements in the field.'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response2"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}