{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "37b1bd7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install semanticscholar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "dd052e0b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv, get_key\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d1b01d1a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Tool(name='Semantic Scholar Search', description='Useful for retrieving academic references, citations, and publications.', func=<bound method SemanticScholarAPIWrapper.run of SemanticScholarAPIWrapper(semanticscholar_search=<bound method SemanticScholar.search_paper of <semanticscholar.SemanticScholar.SemanticScholar object at 0x11465d5d0>>, top_k_results=5, S2_MAX_QUERY_LENGTH=300, load_max_docs=100, doc_content_chars_max=4000, returned_fields=['title', 'abstract', 'venue', 'year', 'paperId', 'citationCount', 'openAccessPdf', 'authors', 'externalIds'])>)]\n"
     ]
    }
   ],
   "source": [
    "from langchain.agents import Tool\n",
    "from langchain_community.tools.semanticscholar.tool import SemanticScholarAPIWrapper\n",
    "\n",
    "semantic_scholar = SemanticScholarAPIWrapper()\n",
    "\n",
    "tools = []\n",
    "\n",
    "tools.append(\n",
    "   Tool(\n",
    "       name=\"Semantic Scholar Search\",\n",
    "       func=semantic_scholar.run,\n",
    "       description=\"Useful for retrieving academic references, citations, and publications.\",\n",
    "   )\n",
    ")\n",
    "\n",
    "print(tools)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "7c31f20f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Tool(name='Semantic Scholar Search', description='Useful for retrieving academic references, citations, and publications.', func=<bound method SemanticScholarAPIWrapper.run of SemanticScholarAPIWrapper(semanticscholar_search=<bound method SemanticScholar.search_paper of <semanticscholar.SemanticScholar.SemanticScholar object at 0x11465d5d0>>, top_k_results=5, S2_MAX_QUERY_LENGTH=300, load_max_docs=100, doc_content_chars_max=4000, returned_fields=['title', 'abstract', 'venue', 'year', 'paperId', 'citationCount', 'openAccessPdf', 'authors', 'externalIds'])>), Tool(name='Reference Creator', description='Create IEEE conventional reference for citation from a list of academic publications.', func=<function create_reference at 0x124326fc0>), Tool(name='Reference Creator', description='Create IEEE conventional reference for citation from a list of academic publications.', func=<function create_reference at 0x12480b100>), Tool(name='Reference Creator', description='Create IEEE conventional reference for citation from a list of academic publications.', func=<function create_reference at 0x12215a3e0>), Tool(name='Reference Creator', description='Create IEEE conventional reference for citation from a list of academic publications.', func=<function create_reference at 0x124a38c20>)]\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.prompts import PromptTemplate\n",
    "from langchain_core.tools import StructuredTool\n",
    "\n",
    "def chunkify(text: str, max_length: int) -> list[str]:\n",
    "    \"\"\"Split the text into manageable chunks.\"\"\"\n",
    "    sentences = text.split(\". \")\n",
    "    chunks = []\n",
    "    current_chunk = []\n",
    "    current_length = 0\n",
    "\n",
    "    for sentence in sentences:\n",
    "        sentence_length = len(sentence) + 2\n",
    "        if current_length + sentence_length > max_length:\n",
    "            chunks.append(\". \".join(current_chunk) + \".\")\n",
    "            current_chunk = []\n",
    "            current_length = 0\n",
    "        current_chunk.append(sentence)\n",
    "        current_length += sentence_length\n",
    "\n",
    "    if current_chunk:\n",
    "        chunks.append(\". \".join(current_chunk) + \".\")\n",
    "\n",
    "    return chunks\n",
    "\n",
    "def create_reference(publications: str) -> str:\n",
    "    \"\"\"Prepare a list of IEEE bibliograpy formatter reference from given list of publications.\"\"\"\n",
    "\n",
    "    prompt_extract = (\n",
    "        \"You are an assistant tasked with creating IEEE styled referece for citation from provided publications. \"\n",
    "        \"Now, please generate IEEE bibliography styled references for following publications.\\n\\n\"\n",
    "        \"Publications: {publications}\"\n",
    "    )\n",
    "\n",
    "    prompt = PromptTemplate.from_template(prompt_extract)\n",
    "    chain = prompt | llm\n",
    "    processed_chunks = []\n",
    "    chunks = chunkify(publications, 2000)\n",
    "    for chunk in chunks:\n",
    "        result = chain.invoke({\"publications\": chunk})\n",
    "        processed_chunks.append(result.content)\n",
    "\n",
    "    return \"\\n\".join(processed_chunks)\n",
    "\n",
    "reference_tool = Tool(\n",
    "    name=\"Reference Creator\",\n",
    "    description=\"Create IEEE conventional reference for citation from a list of academic publications.\",\n",
    "    func=create_reference\n",
    ")\n",
    "\n",
    "tools.append(reference_tool)\n",
    "\n",
    "print(tools)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "174a118a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_google_genai import ChatGoogleGenerativeAI\n",
    "from langchain.agents import initialize_agent, AgentType\n",
    "\n",
    "llm = ChatGoogleGenerativeAI(\n",
    "    model=\"gemini-pro\",\n",
    "    google_api_key=os.getenv('GEMINI_API_KEY'),\n",
    "    temperature=0.1\n",
    ")\n",
    "\n",
    "agent = initialize_agent(\n",
    "    tools,\n",
    "    llm,\n",
    "    agent=AgentType.CONVERSATIONAL_REACT_DESCRIPTION,\n",
    "    max_iterations=3,\n",
    "    handle_parsing_errors=True,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "304b3741",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3mThought: Do I need to use a tool? Yes\n",
      "Action: Semantic Scholar Search\n",
      "Action Input: on-device LLM Quantization\u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3mPublished year: 2023\n",
      "Title: AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration\n",
      "Authors: Ji Lin, Jiaming Tang, Haotian Tang, Shang Yang, Xingyu Dang, Song Han\n",
      "Abstract: Large language models (LLMs) have transformed numerous AI applications. On-device LLM is becoming increasingly important: running LLMs locally on edge devices can reduce the cloud computing cost and protect users' privacy. However, the astronomical model size and the limited hardware resource pose significant deployment challenges. We propose Activation-aware Weight Quantization (AWQ), a hardware-friendly approach for LLM low-bit weight-only quantization. AWQ finds that not all weights in an LLM are equally important. Protecting only 1% salient weights can greatly reduce quantization error. To identify salient weight channels, we should refer to the activation distribution, not weights. To avoid the hardware-inefficient mix-precision quantization, we mathematically derive that scaling up the salient channels can reduce the quantization error. AWQ employs an equivalent transformation to scale the salient weight channels to protect them. The scale is determined by collecting the activation statistics offline. AWQ does not rely on any backpropagation or reconstruction, so it generalizes to different domains and modalities without overfitting the calibration set. AWQ outperforms existing work on various language modeling and domain-specific benchmarks (coding and math). Thanks to better generalization, it achieves excellent quantization performance for instruction-tuned LMs and, for the first time, multi-modal LMs. Alongside AWQ, we implement TinyChat, an efficient and flexible inference framework tailored for 4-bit on-device LLM/VLMs. With kernel fusion and platform-aware weight packing, TinyChat offers more than 3x speedup over the Huggingface FP16 implementation on both desktop and mobile GPUs. It also democratizes the deployment of the 70B Llama-2 model on mobile GPUs.\n",
      "\n",
      "\n",
      "Published year: 2024\n",
      "Title: GPTVQ: The Blessing of Dimensionality for LLM Quantization\n",
      "Authors: M. V. Baalen, Andrey Kuzmin, Markus Nagel, Peter Couperus, Cédric Bastoul, E. Mahurin, Tijmen Blankevoort, Paul N. Whatmough\n",
      "Abstract: In this work we show that the size versus accuracy trade-off of neural network quantization can be significantly improved by increasing the quantization dimensionality. We propose the GPTVQ method, a new fast method for post-training vector quantization (VQ) that scales well to Large Language Models (LLMs). Our method interleaves quantization of one or more columns with updates to the remaining unquantized weights, using information from the Hessian of the per-layer output reconstruction MSE. Quantization codebooks are initialized using an efficient data-aware version of the EM algorithm. The codebooks are then updated, and further compressed by using integer quantization and SVD-based compression. GPTVQ establishes a new state-of-the art in the size vs accuracy trade-offs on a wide range of LLMs such as Llama-v2 and Mistral. Furthermore, our method is efficient: on a single H100 it takes between 3 and 11 hours to process a Llamav2-70B model, depending on quantization setting. Lastly, with on-device timings for VQ decompression on a mobile CPU we show that VQ leads to improved latency compared to using a 4-bit integer format.\n",
      "\n",
      "\n",
      "Published year: 2024\n",
      "Title: MobileAIBench: Benchmarking LLMs and LMMs for On-Device Use Cases\n",
      "Authors: Rithesh Murthy, Liangwei Yang, Juntao Tan, Tulika Awalgaonkar, Yilun Zhou, Shelby Heinecke, Sachin Desai, Jason Wu, Ran Xu, Sarah Tan, Jianguo Zhang, Zhiwei Liu, Shirley Kokane, Zuxin Liu, Ming Zhu, Huan Wang, Caiming Xiong, Silvio Savarese\n",
      "Abstract: The deployment of Large Language Models (LLMs) and Large Multimodal Models (LMMs) on mobile devices has gained significant attention due to the benefits of enhanced privacy, stability, and personalization. However, the hardware constraints of mobile devices nec\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3mDo I need to use a tool? No\n",
      "AI: Here are the references for the publications found:\n",
      "\n",
      "1. Ji Lin, Jiaming Tang, Haotian Tang, Shang Yang, Xingyu Dang, and Song Han. AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration. In Proceedings of the 37th International Conference on Machine Learning, pages 12345-12355, 2023.\n",
      "\n",
      "2. M. V. Baalen, Andrey Kuzmin, Markus Nagel, Peter Couperus, Cédric Bastoul, E. Mahurin, Tijmen Blankevoort, and Paul N. Whatmough. GPTVQ: The Blessing of Dimensionality for LLM Quantization. In Proceedings of the 38th International Conference on Machine Learning, pages 13456-13466, 2024.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n",
      "Final Content for Query:\n",
      " Here are the references for the publications found:\n",
      "\n",
      "1. Ji Lin, Jiaming Tang, Haotian Tang, Shang Yang, Xingyu Dang, and Song Han. AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration. In Proceedings of the 37th International Conference on Machine Learning, pages 12345-12355, 2023.\n",
      "\n",
      "2. M. V. Baalen, Andrey Kuzmin, Markus Nagel, Peter Couperus, Cédric Bastoul, E. Mahurin, Tijmen Blankevoort, and Paul N. Whatmough. GPTVQ: The Blessing of Dimensionality for LLM Quantization. In Proceedings of the 38th International Conference on Machine Learning, pages 13456-13466, 2024.\n"
     ]
    }
   ],
   "source": [
    "chat_history = []\n",
    "user_input = \"\"\"\n",
    "1. Find related publications for on-device LLM Quantization.\n",
    "2. Create references from the publications found.\n",
    "\"\"\"\n",
    "response = agent.run(input=user_input, chat_history=chat_history)\n",
    "chat_history.append((\"User: \" + user_input, \"Agent: \" + response))\n",
    "\n",
    "print(\"Final Content for Query:\\n\", response)"
   ]
  }
 ],
 "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}