rag_inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RAG Inference\n",
    "\n",
    "This notebook follows `rag_create_index.ipynb` where we set up a RAG index. Here we actually use that index to interact with RAG-enabled cosmosage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load model, tokenizer, retriever\n",
    "from langchain.vectorstores import FAISS\n",
    "from langchain.embeddings import HuggingFaceEmbeddings\n",
    "import torch\n",
    "from auto_gptq import AutoGPTQForCausalLM\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "\n",
    "model_dir = \"models/cosmosage_v2/\"\n",
    "#model = AutoGPTQForCausalLM.from_quantized(model_dir) # not yet supported\n",
    "model = AutoModelForCausalLM.from_pretrained(model_dir, torch_dtype=torch.bfloat16).to('cuda')\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_dir)\n",
    "embeddings = HuggingFaceEmbeddings(model_name='BAAI/bge-large-en-v1.5')\n",
    "index = FAISS.load_local(f\"datasets/faiss_index_arxiv_long\", embeddings)\n",
    "retriever = index.as_retriever(search_type=\"mmr\", search_kwargs={'k': 4})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.llms import HuggingFacePipeline\n",
    "from langchain.prompts import PromptTemplate\n",
    "from transformers import pipeline\n",
    "from langchain.chains import LLMChain\n",
    "from langchain.schema.runnable import RunnablePassthrough\n",
    "\n",
    "text_generation_pipeline = pipeline(\n",
    "    model=model,\n",
    "    tokenizer=tokenizer,\n",
    "    task=\"text-generation\",\n",
    "    temperature=0.7,\n",
    "    repetition_penalty=1.01,\n",
    "    return_full_text=True,\n",
    "    max_new_tokens=1000,\n",
    "    do_sample=True,\n",
    ")\n",
    "\n",
    "llm = HuggingFacePipeline(pipeline=text_generation_pipeline)\n",
    "\n",
    "prompt_template = \"\"\"You are cosmosage, an AI programmed to provide excellent and detailed answers to the user's question. You are an expert cosmology assistant, able to answer questions on the cosmic microwave background, galaxy formation, large scale structure, theoretical cosmology, inflation, big bang nucleosynthesis, cosmology instrumentation, and other related topics. Please assume the user is fluent in scientific terminology. Elaborate where possible to give a complete answer. If you do not know, say you do not know. \n",
    "SUPPLEMENTAL INFORMATION: {context}▁ USER: {question}▁ ASSISTANT:\"\"\"\n",
    "\n",
    "prompt = PromptTemplate(input_variables=[\"context\", \"question\"], template=prompt_template)\n",
    "llm_chain = LLMChain(llm=llm, prompt=prompt)\n",
    "rag_chain = (\n",
    " {\"context\": retriever, \"question\": RunnablePassthrough()}\n",
    "    | llm_chain\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"How many parameters does the flat wCDM model have?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = llm_chain.invoke({\"context\":\"\", \"question\": question})\n",
    "print(result['text'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result_rag = rag_chain.invoke(question)\n",
    "print(result_rag['text'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result_rag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}