rag.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GolfGPT RAG system"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Objective:** Letting golfers get fast and accurate rule answers to their on-course questions.\n",
    "\n",
    "\n",
    "**Structure**\n",
    "1. Imports (packages & data)\n",
    "2. Vector embeddings: Data loading, chunking, embedding and vector store injection\n",
    "3. Model setup (model instantiation, incl. hyperparameters)\n",
    "4. Prompt template preparation (incl. system prompt, chat history and query parameters)\n",
    "5. Interaction logic (receive user input, convert to embeddings, retrieval logic, answer generation)\n",
    "6. Deployment (gradio / hf spaces)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/miniconda3/envs/golfgpt/lib/python3.12/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"
     ]
    }
   ],
   "source": [
    "# Importing necessary libraries\n",
    "from dotenv import load_dotenv\n",
    "from typing import List\n",
    "from langchain_openai import OpenAIEmbeddings, ChatOpenAI\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from langchain_community.vectorstores import Chroma\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.runnables import RunnableParallel\n",
    "import gradio as gr\n",
    "load_dotenv();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the rules\n",
    "with open('rules.txt', 'r') as file:\n",
    "    golf_rules = file.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created 25 major rule chunks\n",
      "Created 379 total chunks\n"
     ]
    }
   ],
   "source": [
    "# First split on major rule boundaries\n",
    "major_splitter = RecursiveCharacterTextSplitter(\n",
    "    separators=[r\"\\n\\*\\*\\*\\nRule\"],\n",
    "    chunk_size=10000,  # Large enough to capture entire rules\n",
    "    chunk_overlap=0,\n",
    "    length_function=len,\n",
    "    is_separator_regex=True,\n",
    ")\n",
    "\n",
    "# Then split large chunks into smaller ones\n",
    "detail_splitter = RecursiveCharacterTextSplitter(\n",
    "    chunk_size=1000,\n",
    "    chunk_overlap=200,\n",
    "    length_function=len,\n",
    ")\n",
    "\n",
    "# First split into major rules\n",
    "major_chunks = major_splitter.split_text(golf_rules)\n",
    "print(f\"Created {len(major_chunks)} major rule chunks\")\n",
    "\n",
    "# Then split large chunks into smaller ones\n",
    "chunks = []\n",
    "for chunk in major_chunks:\n",
    "    if len(chunk) > 1000:\n",
    "        sub_chunks = detail_splitter.split_text(chunk)\n",
    "        chunks.extend(sub_chunks)\n",
    "    else:\n",
    "        chunks.append(chunk)\n",
    "\n",
    "print(f\"Created {len(chunks)} total chunks\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiating embeddings model and LLM\n",
    "embeddings = OpenAIEmbeddings()\n",
    "llm = ChatOpenAI(temperature=0, model=\"gpt-4o-mini\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create vector store\n",
    "vectorstore = Chroma.from_texts(\n",
    "    texts=major_chunks,\n",
    "    embedding=embeddings,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create prompt template\n",
    "template = \"\"\"You are a helpful golf rules assistant. Use the following pieces of context to answer the question at the end. \n",
    "If you don't know the answer, just say that you don't know, don't try to make up an answer. \n",
    "You can only answer questions about the rules of golf. If a question is not about golf, kindly remind them that you only are a golf rules assistant.\n",
    "Think step by step and remember to use emojis and cheer the golfer on!\n",
    "\n",
    "Context: {context}\n",
    "\n",
    "Question: {question}\n",
    "\n",
    "Answer:\"\"\"\n",
    "\n",
    "prompt = ChatPromptTemplate.from_template(template)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create RAG chain\n",
    "def format_docs(docs):\n",
    "    formatted_docs = []\n",
    "    for i, doc in enumerate(docs, 1):\n",
    "        formatted_docs.append(f\"[Source {i}]: {doc.page_content}\")\n",
    "    return \"\\n\\n\".join(formatted_docs)\n",
    "\n",
    "def format_response(response, doctitle):\n",
    "        return f\"{response}\\n\\n{'='*36}\\nSource used: {doctitle}\"\n",
    "\n",
    "retriever = vectorstore.as_retriever(search_kwargs={\"k\": 1})\n",
    "\n",
    "def rag_chain_with_sources(question):\n",
    "    docs = retriever.invoke(question)\n",
    "    chain = (\n",
    "        RunnableParallel({\n",
    "            \"context\": lambda _: format_docs(docs),\n",
    "            \"question\": lambda _: question\n",
    "        })\n",
    "        | prompt\n",
    "        | llm\n",
    "        | StrOutputParser()\n",
    "    )\n",
    "    response = chain.invoke({})\n",
    "    return response, docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define function to query the RAG system\n",
    "def query_golf_rules(question: str) -> str:\n",
    "    response, docs = rag_chain_with_sources(question)\n",
    "    content_lines = [line for line in docs[0].page_content.split(\"\\n\") if line.strip() and line.strip() != \"***\"]\n",
    "    doctitle = content_lines[0] if content_lines else \"Unknown Rule\"\n",
    "    \n",
    "    return format_response(response, doctitle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "* Running on local URL:  http://127.0.0.1:7861\n",
      "* Running on public URL: https://6891e12fb288757636.gradio.live\n",
      "\n",
      "This share link expires in 72 hours. For free permanent hosting and GPU upgrades, run `gradio deploy` from the terminal in the working directory to deploy to Hugging Face Spaces (https://huggingface.co/spaces)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"https://6891e12fb288757636.gradio.live\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Configure Gradio interface\n",
    "def gradio_interface(question):\n",
    "    return query_golf_rules(question)\n",
    "\n",
    "demo = gr.Interface(\n",
    "    fn=gradio_interface,\n",
    "    inputs=gr.Textbox(\n",
    "        lines=2, \n",
    "        placeholder=\"What would you like to know?\",\n",
    "        label=\"Your Question\"\n",
    "    ),\n",
    "    outputs=gr.Textbox(\n",
    "        lines=10,\n",
    "        label=\"GolfGPT Answer\"\n",
    "    ),\n",
    "    title=\"GolfGPT Rules Assistant\",\n",
    "    description=\"Ask questions about golf rules and get accurate answers based on the official rules of golf. The model can make mistakes\",\n",
    "    examples=[\n",
    "        \"What are the rules for taking a drop?\",\n",
    "        \"How do I handle a lost ball?\",\n",
    "        \"Can I repair ball marks on the green?\",\n",
    "        \"What are the rules for playing from a bunker?\",\n",
    "        \"How do I handle an unplayable lie?\"\n",
    "    ],\n",
    "    theme=gr.themes.Soft()\n",
    ")\n",
    "\n",
    "# Launch the interface\n",
    "demo.launch(pwa=True, share=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "golfgpt",
   "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}