Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+data/videos_subset.db/chroma.sqlite3 filter=lfs diff=lfs merge=lfs -text
+data/videos_subset_more_context.db/chroma.sqlite3 filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,138 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/

ETL_Walkthrough.ipynb

@@ -0,0 +1,375 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# ETL to get the text data from the playlist\n",
+    "\n",
+    "This notebook shows the process of building the corpus of transcripts from the YouTube playlist.\n",
+    "\n",
+    "**Extract**: Pull data (transcripts) from each video.  \n",
+    "**Transform**:  \n",
+    "**Load**: Load data into our database where it will be retrieved from.  "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from models import etl\n",
+    "import json"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "First we load the video information. This includes the video IDs and titles."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open('data/single_video.json') as f:\n",
+    "    video_info = json.load(f)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Then we must extract the transcripts using the YouTube Transcript API. This is done over all of the videos.  \n",
+    "This produces a list of video segments with timestamps.  \n",
+    "Next, we format the transcript by adding metadata so that the segments are easily identified for retreival later.  \n",
+    "Since the original segments are small, they are batched with overlap to preserve semantic meaning."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "get_video_transcript took 0.84 seconds.\n",
+      "Transcript for video 5sLYAQS9sWQ fetched.\n"
+     ]
+    }
+   ],
+   "source": [
+    "videos = []\n",
+    "for video in video_info:\n",
+    "    video_id = video[\"id\"]\n",
+    "    video_title = video[\"title\"]\n",
+    "    transcript = etl.get_video_transcript(video_id)\n",
+    "    print(f\"Transcript for video {video_id} fetched.\")\n",
+    "    if transcript:\n",
+    "        formatted_transcript = etl.format_transcript(transcript, video_id, video_title, batch_size=5, overlap=2)\n",
+    "        \n",
+    "        videos.extend(formatted_transcript)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The last step is to load the data into a database. We will use a Chromadb database.  \n",
+    "The embedding function is the ____ model from HuggingFace."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Database created at data/single_video.db\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Initialize the database\n",
+    "from utils.embedding_utils import MyEmbeddingFunction\n",
+    "import chromadb\n",
+    "\n",
+    "embed_text = MyEmbeddingFunction()\n",
+    "\n",
+    "db_path = \"data/single_video.db\"\n",
+    "client = chromadb.PersistentClient(path=db_path)\n",
+    "\n",
+    "client.create_collection(\n",
+    "    name=\"huberman_videos\",\n",
+    "    embedding_function=embed_text,\n",
+    "    metadata={\"hnsw:space\": \"cosine\"}\n",
+    ")\n",
+    "\n",
+    "print(f\"Database created at {db_path}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Data loaded to database at data/single_video.db.\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Add the data to the database\n",
+    "client = chromadb.PersistentClient(path=db_path)\n",
+    "    \n",
+    "collection = client.get_collection(\"huberman_videos\")\n",
+    "\n",
+    "documents = [segment['text'] for segment in videos]\n",
+    "metadata = [segment['metadata'] for segment in videos]\n",
+    "ids = [segment['metadata']['segment_id'] for segment in videos]\n",
+    "\n",
+    "collection.add(\n",
+    "    documents=documents,\n",
+    "    metadatas=metadata,\n",
+    "    ids=ids\n",
+    ")\n",
+    "\n",
+    "print(f\"Data loaded to database at {db_path}.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Here is some of the data:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Number of segments: 26\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>ids</th>\n",
+       "      <th>embeddings</th>\n",
+       "      <th>metadatas</th>\n",
+       "      <th>documents</th>\n",
+       "      <th>uris</th>\n",
+       "      <th>data</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>5sLYAQS9sWQ__0</td>\n",
+       "      <td>[-0.11489544063806534, -0.03262839838862419, -...</td>\n",
+       "      <td>{'segment_id': '5sLYAQS9sWQ__0', 'source': 'ht...</td>\n",
+       "      <td>GPT, or Generative Pre-trained Transformer, is...</td>\n",
+       "      <td>None</td>\n",
+       "      <td>None</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>5sLYAQS9sWQ__12</td>\n",
+       "      <td>[0.094169981777668, -0.10430295020341873, 0.02...</td>\n",
+       "      <td>{'segment_id': '5sLYAQS9sWQ__12', 'source': 'h...</td>\n",
+       "      <td>Now foundation models are pre-trained on large...</td>\n",
+       "      <td>None</td>\n",
+       "      <td>None</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>5sLYAQS9sWQ__15</td>\n",
+       "      <td>[0.042587604373693466, -0.061460819095373154, ...</td>\n",
+       "      <td>{'segment_id': '5sLYAQS9sWQ__15', 'source': 'h...</td>\n",
+       "      <td>I'm talking about things like code. Now, large...</td>\n",
+       "      <td>None</td>\n",
+       "      <td>None</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>5sLYAQS9sWQ__18</td>\n",
+       "      <td>[-0.0245895367115736, -0.058405470103025436, -...</td>\n",
+       "      <td>{'segment_id': '5sLYAQS9sWQ__18', 'source': 'h...</td>\n",
+       "      <td>these models can be tens of gigabytes in size ...</td>\n",
+       "      <td>None</td>\n",
+       "      <td>None</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>5sLYAQS9sWQ__21</td>\n",
+       "      <td>[0.05348338559269905, -0.016104578971862793, -...</td>\n",
+       "      <td>{'segment_id': '5sLYAQS9sWQ__21', 'source': 'h...</td>\n",
+       "      <td>So to put that into perspective, a text file t...</td>\n",
+       "      <td>None</td>\n",
+       "      <td>None</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>5</th>\n",
+       "      <td>5sLYAQS9sWQ__24</td>\n",
+       "      <td>[0.07004527002573013, -0.08996045589447021, -0...</td>\n",
+       "      <td>{'segment_id': '5sLYAQS9sWQ__24', 'source': 'h...</td>\n",
+       "      <td>A lot of words just in one Gb. And how many gi...</td>\n",
+       "      <td>None</td>\n",
+       "      <td>None</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6</th>\n",
+       "      <td>5sLYAQS9sWQ__27</td>\n",
+       "      <td>[0.0283487681299448, -0.11020224541425705, -0....</td>\n",
+       "      <td>{'segment_id': '5sLYAQS9sWQ__27', 'source': 'h...</td>\n",
+       "      <td>Yeah, that's truly a lot of text. And LLMs are...</td>\n",
+       "      <td>None</td>\n",
+       "      <td>None</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>7</th>\n",
+       "      <td>5sLYAQS9sWQ__3</td>\n",
+       "      <td>[-0.0700172707438469, -0.061202701181173325, -...</td>\n",
+       "      <td>{'segment_id': '5sLYAQS9sWQ__3', 'source': 'ht...</td>\n",
+       "      <td>And I've been using GPT in its various forms f...</td>\n",
+       "      <td>None</td>\n",
+       "      <td>None</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>8</th>\n",
+       "      <td>5sLYAQS9sWQ__30</td>\n",
+       "      <td>[-0.04904637485742569, -0.1277533322572708, -0...</td>\n",
+       "      <td>{'segment_id': '5sLYAQS9sWQ__30', 'source': 'h...</td>\n",
+       "      <td>and the more parameters a model has, the more ...</td>\n",
+       "      <td>None</td>\n",
+       "      <td>None</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>9</th>\n",
+       "      <td>5sLYAQS9sWQ__33</td>\n",
+       "      <td>[0.03286760300397873, -0.041724931448698044, 0...</td>\n",
+       "      <td>{'segment_id': '5sLYAQS9sWQ__33', 'source': 'h...</td>\n",
+       "      <td>All right, so how do they work? Well, we can t...</td>\n",
+       "      <td>None</td>\n",
+       "      <td>None</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "               ids                                         embeddings  \\\n",
+       "0   5sLYAQS9sWQ__0  [-0.11489544063806534, -0.03262839838862419, -...   \n",
+       "1  5sLYAQS9sWQ__12  [0.094169981777668, -0.10430295020341873, 0.02...   \n",
+       "2  5sLYAQS9sWQ__15  [0.042587604373693466, -0.061460819095373154, ...   \n",
+       "3  5sLYAQS9sWQ__18  [-0.0245895367115736, -0.058405470103025436, -...   \n",
+       "4  5sLYAQS9sWQ__21  [0.05348338559269905, -0.016104578971862793, -...   \n",
+       "5  5sLYAQS9sWQ__24  [0.07004527002573013, -0.08996045589447021, -0...   \n",
+       "6  5sLYAQS9sWQ__27  [0.0283487681299448, -0.11020224541425705, -0....   \n",
+       "7   5sLYAQS9sWQ__3  [-0.0700172707438469, -0.061202701181173325, -...   \n",
+       "8  5sLYAQS9sWQ__30  [-0.04904637485742569, -0.1277533322572708, -0...   \n",
+       "9  5sLYAQS9sWQ__33  [0.03286760300397873, -0.041724931448698044, 0...   \n",
+       "\n",
+       "                                           metadatas  \\\n",
+       "0  {'segment_id': '5sLYAQS9sWQ__0', 'source': 'ht...   \n",
+       "1  {'segment_id': '5sLYAQS9sWQ__12', 'source': 'h...   \n",
+       "2  {'segment_id': '5sLYAQS9sWQ__15', 'source': 'h...   \n",
+       "3  {'segment_id': '5sLYAQS9sWQ__18', 'source': 'h...   \n",
+       "4  {'segment_id': '5sLYAQS9sWQ__21', 'source': 'h...   \n",
+       "5  {'segment_id': '5sLYAQS9sWQ__24', 'source': 'h...   \n",
+       "6  {'segment_id': '5sLYAQS9sWQ__27', 'source': 'h...   \n",
+       "7  {'segment_id': '5sLYAQS9sWQ__3', 'source': 'ht...   \n",
+       "8  {'segment_id': '5sLYAQS9sWQ__30', 'source': 'h...   \n",
+       "9  {'segment_id': '5sLYAQS9sWQ__33', 'source': 'h...   \n",
+       "\n",
+       "                                           documents  uris  data  \n",
+       "0  GPT, or Generative Pre-trained Transformer, is...  None  None  \n",
+       "1  Now foundation models are pre-trained on large...  None  None  \n",
+       "2  I'm talking about things like code. Now, large...  None  None  \n",
+       "3  these models can be tens of gigabytes in size ...  None  None  \n",
+       "4  So to put that into perspective, a text file t...  None  None  \n",
+       "5  A lot of words just in one Gb. And how many gi...  None  None  \n",
+       "6  Yeah, that's truly a lot of text. And LLMs are...  None  None  \n",
+       "7  And I've been using GPT in its various forms f...  None  None  \n",
+       "8  and the more parameters a model has, the more ...  None  None  \n",
+       "9  All right, so how do they work? Well, we can t...  None  None  "
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import pandas as pd\n",
+    "\n",
+    "client = chromadb.PersistentClient('data/single_video.db')\n",
+    "collection= client.get_collection('huberman_videos')\n",
+    "\n",
+    "num_segments = collection.count()\n",
+    "sample_data = collection.peek()\n",
+    "\n",
+    "transcript_df = pd.DataFrame(sample_data)\n",
+    "\n",
+    "print(f\"Number of segments: {num_segments}\")\n",
+    "transcript_df"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "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.1"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

README.md

@@ -1,12 +1,6 @@
 ---
-title: Fitness QA Bot
-emoji: 🔥
-colorFrom: pink
-colorTo: yellow
+title: Fitness_QA_Bot
+app_file: main.py
 sdk: gradio
-sdk_version: 4.17.0
-app_file: app.py
-pinned: false
+sdk_version: 4.16.0
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

TESTING.ipynb

@@ -0,0 +1,165 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from models import etl"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "etl.main(json_path='data/single_video.json', db='data/single_video.db', batch_size=5, overlap=2)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import chromadb\n",
+    "from models import etl"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "client = chromadb.PersistentClient('data/single_video.db')\n",
+    "collection= client.get_collection('huberman_videos')\n",
+    "# collection.count()\n",
+    "# collection.peek()\n",
+    "\n",
+    "query_text = \"What are the components of an LLM?\"\n",
+    "query_embedding = etl.embed_text(query_text)\n",
+    "results = collection.query(query_texts=[query_text], n_results=5)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'ids': [['5sLYAQS9sWQ__33',\n",
+       "   '5sLYAQS9sWQ__36',\n",
+       "   '5sLYAQS9sWQ__3',\n",
+       "   '5sLYAQS9sWQ__6',\n",
+       "   '5sLYAQS9sWQ__27']],\n",
+       " 'distances': [[0.27329726119651687,\n",
+       "   0.3594438065792097,\n",
+       "   0.4730243492988927,\n",
+       "   0.5004446084705303,\n",
+       "   0.5766584257317211]],\n",
+       " 'metadatas': [[{'segment_id': '5sLYAQS9sWQ__33',\n",
+       "    'source': 'https://www.youtube.com/watch?v=5sLYAQS9sWQ&t=145.328s',\n",
+       "    'title': 'How Large Language Models Work',\n",
+       "    'video_id': '5sLYAQS9sWQ'},\n",
+       "   {'segment_id': '5sLYAQS9sWQ__36',\n",
+       "    'source': 'https://www.youtube.com/watch?v=5sLYAQS9sWQ&t=154.367s',\n",
+       "    'title': 'How Large Language Models Work',\n",
+       "    'video_id': '5sLYAQS9sWQ'},\n",
+       "   {'segment_id': '5sLYAQS9sWQ__3',\n",
+       "    'source': 'https://www.youtube.com/watch?v=5sLYAQS9sWQ&t=10.783s',\n",
+       "    'title': 'How Large Language Models Work',\n",
+       "    'video_id': '5sLYAQS9sWQ'},\n",
+       "   {'segment_id': '5sLYAQS9sWQ__6',\n",
+       "    'source': 'https://www.youtube.com/watch?v=5sLYAQS9sWQ&t=22.544s',\n",
+       "    'title': 'How Large Language Models Work',\n",
+       "    'video_id': '5sLYAQS9sWQ'},\n",
+       "   {'segment_id': '5sLYAQS9sWQ__27',\n",
+       "    'source': 'https://www.youtube.com/watch?v=5sLYAQS9sWQ&t=117.572s',\n",
+       "    'title': 'How Large Language Models Work',\n",
+       "    'video_id': '5sLYAQS9sWQ'}]],\n",
+       " 'embeddings': None,\n",
+       " 'documents': [['All right, so how do they work? Well, we can think of it like this. LLM equals three things: data, architecture, and lastly, we can think of it as training. Those three things are really the components of an LLM.',\n",
+       "   \"data, architecture, and lastly, we can think of it as training. Those three things are really the components of an LLM. Now, we've already discussed the enormous amounts of text data that goes into these things. As for the architecture, this is a neural network and for GPT that is a transformer.\",\n",
+       "   'And I\\'ve been using GPT in its various forms for years. In this video we are going to number 1, ask \"what is an LLM?\" Number 2, we are going to describe how they work. And then number 3,',\n",
+       "   'Number 2, we are going to describe how they work. And then number 3, we\\'re going to ask, \"what are the business applications of LLMs?\" So let\\'s start with number 1, \"what is a large language model?\" Well, a large language model',\n",
+       "   \"Yeah, that's truly a lot of text. And LLMs are also among the biggest models when it comes to parameter count. A parameter is a value the model can change independently as it learns, and the more parameters a model has, the more complex it can be. GPT-3, for example, is pre-trained on a corpus of actually 45 terabytes of data,\"]],\n",
+       " 'uris': None,\n",
+       " 'data': None}"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CONTEXT: All right, so how do they work? Well, we can think of it like this. LLM equals three things: data, architecture, and lastly, we can think of it as training. Those three things are really the components of an LLM.\n",
+      "TITLE: How Large Language Models Work\n",
+      "SOURCE: https://www.youtube.com/watch?v=5sLYAQS9sWQ&t=145.328s\n",
+      "\n",
+      "CONTEXT: data, architecture, and lastly, we can think of it as training. Those three things are really the components of an LLM. Now, we've already discussed the enormous amounts of text data that goes into these things. As for the architecture, this is a neural network and for GPT that is a transformer.\n",
+      "TITLE: How Large Language Models Work\n",
+      "SOURCE: https://www.youtube.com/watch?v=5sLYAQS9sWQ&t=154.367s\n",
+      "\n",
+      "CONTEXT: And I've been using GPT in its various forms for years. In this video we are going to number 1, ask \"what is an LLM?\" Number 2, we are going to describe how they work. And then number 3,\n",
+      "TITLE: How Large Language Models Work\n",
+      "SOURCE: https://www.youtube.com/watch?v=5sLYAQS9sWQ&t=10.783s\n",
+      "\n",
+      "CONTEXT: Number 2, we are going to describe how they work. And then number 3, we're going to ask, \"what are the business applications of LLMs?\" So let's start with number 1, \"what is a large language model?\" Well, a large language model\n",
+      "TITLE: How Large Language Models Work\n",
+      "SOURCE: https://www.youtube.com/watch?v=5sLYAQS9sWQ&t=22.544s\n",
+      "\n",
+      "CONTEXT: Yeah, that's truly a lot of text. And LLMs are also among the biggest models when it comes to parameter count. A parameter is a value the model can change independently as it learns, and the more parameters a model has, the more complex it can be. GPT-3, for example, is pre-trained on a corpus of actually 45 terabytes of data,\n",
+      "TITLE: How Large Language Models Work\n",
+      "SOURCE: https://www.youtube.com/watch?v=5sLYAQS9sWQ&t=117.572s\n",
+      "\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from models.llm import format_context\n",
+    "\n",
+    "print(format_context(results))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "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.1"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

data/logs/videos_subset_more_context_load_log.json

@@ -0,0 +1 @@
+{"videos_info_path": "data/videos_subset.json", "db_path": "data/videos_subset_more_context.db", "batch_size": 15, "overlap": 10, "load_time": "2024-02-07 19:44:01.376335"}

data/single_video.db/2b08fbf8-a94d-47d9-97aa-102f2895fbd0/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d3c9fd302f000d7790aa403c2d0d8fec363fe46f30b07d53020b6e33b22435a9
+size 1676000

data/single_video.db/2b08fbf8-a94d-47d9-97aa-102f2895fbd0/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e87a1dc8bcae6f2c4bea6d5dd5005454d4dace8637dae29bff3c037ea771411e
+size 100

data/single_video.db/2b08fbf8-a94d-47d9-97aa-102f2895fbd0/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:135da168ae29fd55e9cbd2f05414e81b32cd0783d38340b404e315cebd7a7647
+size 4000

data/single_video.json

@@ -0,0 +1,6 @@
+[
+    {
+        "id": "5sLYAQS9sWQ",
+        "title": "How Large Language Models Work"
+    }
+]

data/videos.json

@@ -0,0 +1,70 @@
+[
+    {
+        "id": "LYYyQcAJZfk",
+        "title": "Science-Supported Tools to Accelerate Your Fitness Goals | Huberman Lab Podcast"
+    },
+    {
+        "id": "q37ARYnRDGc",
+        "title": "Dr. Andy Galpin: Optimal Nutrition & Supplementation for Fitness | Huberman Lab Guest Series"
+    },
+    {
+        "id": "juD99_sPWGU",
+        "title": "Dr. Andy Galpin: Maximize Recovery to Achieve Fitness & Performance Goals | Huberman Lab"
+    },
+    {
+        "id": "UIy-WQCZd4M",
+        "title": "Dr. Andy Galpin: Optimize Your Training Program for Fitness & Longevity | Huberman Lab Guest Series"
+    },
+    {
+        "id": "oNkDA2F7CjM",
+        "title": "Dr. Andy Galpin: How to Build Physical Endurance & Lose Fat | Huberman Lab Guest Series"
+    },
+    {
+        "id": "CyDLbrZK75U",
+        "title": "Dr. Andy Galpin: Optimal Protocols to Build Strength & Grow Muscles | Huberman Lab Guest Series"
+    },
+    {
+        "id": "zEYE-vcVKy8",
+        "title": "Dr. Andy Galpin: How to Assess & Improve All Aspects of Your Fitness | Huberman Lab Guest Series"
+    },
+    {
+        "id": "q1Ss8sTbFBY",
+        "title": "Fitness Toolkit: Protocol & Tools to Optimize Physical Health | Huberman Lab Podcast #94"
+    },
+    {
+        "id": "DTCmprPCDqc",
+        "title": "Dr. Peter Attia: Exercise, Nutrition, Hormones for Vitality & Longevity | Huberman Lab Podcast #85"
+    },
+    {
+        "id": "UNCwdFxPtE8",
+        "title": "Jeff Cavaliere: Optimize Your Exercise Program with Science-Based Tools | Huberman Lab Podcast #79"
+    },
+    {
+        "id": "a9yFKPmPZ90",
+        "title": "Ido Portal: The Science & Practice of Movement | Huberman Lab Podcast #77"
+    },
+    {
+        "id": "tkH2-_jMCSk",
+        "title": "Improve Flexibility with Research-Supported Stretching Protocols | Huberman Lab Podcast #76"
+    },
+    {
+        "id": "IAnhFUUCq6c",
+        "title": "Dr. Andy Galpin: How to Build Strength, Muscle Size & Endurance | Huberman Lab Podcast #65"
+    },
+    {
+        "id": "GLgKkG44MGo",
+        "title": "Dr. Jack Feldman: Breathing for Mental & Physical Health & Performance | Huberman Lab Podcast #54"
+    },
+    {
+        "id": "iMvtHqLmEkI",
+        "title": "Dr. Duncan French: How to Exercise for Strength Gains & Hormone Optimization | Huberman Lab #45"
+    },
+    {
+        "id": "VQLU7gpk_X8",
+        "title": "How to Build Endurance in Your Brain & Body"
+    },
+    {
+        "id": "XLr2RKoD-oY",
+        "title": "Science of Muscle Growth, Increasing Strength & Muscular Recovery"
+    }
+]

data/videos_subset.db/5c5462a4-1205-4d8f-8261-eccc02fe61b4/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e1dc504103bbe12bbc64b0a140e9c7eac03ceb4ad42d6b251e8399e6035b6407
+size 3352000

data/videos_subset.db/5c5462a4-1205-4d8f-8261-eccc02fe61b4/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:af2a3cca693212a055b68369ac5fd6015123fe35f04aafb696bf97a944142a76
+size 100

data/videos_subset.db/5c5462a4-1205-4d8f-8261-eccc02fe61b4/index_metadata.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c76fc6a10e46df31db8a10e076a256a53418217dca6d9d7833623c4c1fcbee63
+size 75509

data/videos_subset.db/5c5462a4-1205-4d8f-8261-eccc02fe61b4/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6ad502e522081d0f08a17da3b8d9af56c4a78798cd4b76926924a473e7f210fb
+size 8000

data/videos_subset.db/5c5462a4-1205-4d8f-8261-eccc02fe61b4/link_lists.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:858c69984324e9f15a6fbd27b9f4efdd2ec55fe7240c7c438f39a2d8cbd808fb
+size 18268

data/videos_subset.db/chroma.sqlite3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d6b150b1e85be961e1bcb3b67fa9ded5a20b79bb53ee59d05f5f2b6c09d7501e
+size 13717504

data/videos_subset.json

@@ -0,0 +1,14 @@
+[
+    {
+        "id": "LYYyQcAJZfk",
+        "title": "Science-Supported Tools to Accelerate Your Fitness Goals | Huberman Lab Podcast"
+    },
+    {
+        "id": "q37ARYnRDGc",
+        "title": "Dr. Andy Galpin: Optimal Nutrition & Supplementation for Fitness | Huberman Lab Guest Series"
+    },
+    {
+        "id": "oNkDA2F7CjM",
+        "title": "Dr. Andy Galpin: How to Build Physical Endurance & Lose Fat | Huberman Lab Guest Series"
+    }
+]

data/videos_subset_more_context.db/7dde0ca4-4e76-4236-abfc-198f65e02dc0/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:31040549c44b1b419d5f60f849fd1abd006c4dd33bf077cf498e92360a7a890f
+size 3352000

data/videos_subset_more_context.db/7dde0ca4-4e76-4236-abfc-198f65e02dc0/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:092185a3873be9e8dcabd69df1976197a0afb6cf7d175f3ee73ccae7283d0abc
+size 100

data/videos_subset_more_context.db/7dde0ca4-4e76-4236-abfc-198f65e02dc0/index_metadata.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f2bc9be2510d046cef8a9e0aabc77141b639c086811cdbd4203b15977c7b2f6e
+size 75323

data/videos_subset_more_context.db/7dde0ca4-4e76-4236-abfc-198f65e02dc0/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6cd4536124c25ea26dc53deadb45f017008f2815fbe38deb6e22ae1ecec8e57a
+size 8000

data/videos_subset_more_context.db/7dde0ca4-4e76-4236-abfc-198f65e02dc0/link_lists.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f0e5fb78cc60a74d205b2b5797f65b15ce3ecd5d63d72f4ef82daf716e787a39
+size 18200

data/videos_subset_more_context.db/chroma.sqlite3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bfa8d1ef4a95be2ae4714fbb86e740c4ed4bd334aa068fb64817f84e25857bda
+size 21721088

main.py

@@ -0,0 +1,25 @@
+from models import llm, retrieval
+import gradio as gr
+
+
+def run_query(question, db_path="data/videos_subset_more_context.db", #dev path
+              num_rel_segments=5, 
+              llm_model="gpt-3.5-turbo-0125", 
+              llm_temp=0.1):
+    
+    relevant_segments = retrieval.get_relevant_segments(question,
+                                                         db_path=db_path,
+                                                         n_results=num_rel_segments)
+    
+    answer = llm.answer_with_context(question, 
+                                     relevant_segments, 
+                                     model=llm_model, 
+                                     temperature=llm_temp)
+    
+    return answer
+
+
+if __name__ == "__main__":
+    demo = gr.Interface(fn=run_query, inputs="text", outputs="text")
+    
+    demo.launch(share=True)

models/etl.py

@@ -0,0 +1,126 @@
+import json
+import chromadb
+from datetime import datetime
+
+from utils.general_utils import timeit
+from utils.embedding_utils import MyEmbeddingFunction
+from youtube_transcript_api import YouTubeTranscriptApi
+
+
+@timeit
+def run_etl(json_path="data/videos.json", db=None, batch_size=None, overlap=None):
+    with open(json_path) as f:
+        video_info = json.load(f)
+        
+    videos = []
+    for video in video_info:
+        video_id = video["id"]
+        video_title = video["title"]
+        transcript = get_video_transcript(video_id)
+        print(f"Transcript for video {video_id} fetched.")
+        if transcript:
+            formatted_transcript = format_transcript(transcript, video_id, video_title, batch_size=batch_size, overlap=overlap)
+            
+            videos.extend(formatted_transcript)
+    
+    if db:
+        initialize_db(db)
+        load_data_to_db(db, videos)
+        log_data_load(json_path, db, batch_size, overlap)
+    else:
+        print("No database specified. Skipping database load.")
+        print(videos)
+        
+
+@timeit
+def get_video_transcript(video_id):
+    try:
+        transcript = YouTubeTranscriptApi.get_transcript(video_id, languages=['en', 'en-US'])
+        return transcript
+    except Exception as e:
+        print(f"Error fetching transcript for video {video_id}: {str(e)}")
+        return None
+    
+    
+def format_transcript(transcript, video_id, video_title, batch_size=None, overlap=None):
+    formatted_data = []
+    base_url = f"https://www.youtube.com/watch?v={video_id}"
+    query_params = "&t={start}s"
+    
+    if not batch_size:
+        batch_size = 1
+        overlap = 0
+        
+    for i in range(0, len(transcript), batch_size - overlap):
+        batch = list(transcript[i:i+batch_size])
+        
+        start_time = batch[0]["start"]
+        
+        text = " ".join(entry["text"] for entry in batch)
+        
+        url = base_url + query_params.format(start=start_time)
+        
+        metadata = {
+            "video_id": video_id,
+            "segment_id": video_id + "__" + str(i),
+            "title": video_title,
+            "source": url
+        }
+        
+        segment = {"text": text, "metadata": metadata}
+        
+        formatted_data.append(segment)
+
+    return formatted_data
+
+
+embed_text = MyEmbeddingFunction()
+
+def initialize_db(db_path, distance_metric="cosine"):
+    client = chromadb.PersistentClient(path=db_path)
+    
+    # Clear existing data
+    # client.reset()
+    
+    client.create_collection(
+        name="huberman_videos",
+        embedding_function=embed_text,
+        metadata={"hnsw:space": distance_metric}
+    )
+    
+    print(f"Database created at {db_path}")
+    
+
+def load_data_to_db(db_path, data):
+    client = chromadb.PersistentClient(path=db_path)
+    
+    collection = client.get_collection("huberman_videos")
+
+    documents = [segment['text'] for segment in data]
+    metadata = [segment['metadata'] for segment in data]
+    ids = [segment['metadata']['segment_id'] for segment in data]
+
+    collection.add(
+        documents=documents,
+        metadatas=metadata,
+        ids=ids
+    )
+
+    print(f"Data loaded to database at {db_path}.")
+    
+    
+def log_data_load(json_path, db_path, batch_size, overlap):
+    log_json = json.dumps({
+        "videos_info_path": json_path,
+        "db_path": db_path,
+        "batch_size": batch_size,
+        "overlap": overlap,
+        "load_time": str(datetime.now())
+    })
+    
+    db_file = db_path.split("/")[-1]
+    db_name = db_file.split(".")[0]
+    log_path = f"data/logs/{db_name}_load_log.json"
+    
+    with open(log_path, "w") as f:
+        f.write(log_json)

models/llm.py

@@ -0,0 +1,50 @@
+import openai
+import os
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv()
+
+# Get the API key from the environment variable
+api_key = os.getenv("OPENAI_API_KEY")
+
+# Create the OpenAI client
+client = openai.OpenAI(api_key=api_key)
+
+context_result_base = "CONTEXT: {text}\nTITLE: {title}\nSOURCE: {source}\n\n"
+
+def format_context(db_query_results):
+    documents = db_query_results['documents'][0]
+    metadatas = db_query_results['metadatas'][0]
+    
+    formatted_context = ""
+    for i in range(len(documents)):
+        result_text = context_result_base.format(text=documents[i], title=metadatas[i]['title'], source=metadatas[i]['source'])
+        formatted_context += result_text
+    
+    return formatted_context
+
+
+def answer_with_context(question, context, model="gpt-3.5-turbo-1106", temperature=0.5):
+    formatted_context = format_context(context)
+    
+    instruction = '''You are a question-answering bot. The user will ask a question about fitness and recovery. First, you will be provided relevant context. The relevant context are segments of transcripts from Andrew Huberman's playlist on fitness and recovery where he has conversations about these topics. Answer the user's question and include the video title and link to the relevant context where they talk about the topic of the user's question.  When referencing relevant context, return its TITLE and SOURCE. If no context are related to the question, answer the question yourself and state that "No relevant clips were found". Use this format:
+User: ```What is muscle atrophy?```
+AI: ```Muscle atrophy is the decrease in size and wasting of muscle tissue.
+VIDEO: Example video title
+SOURCE: Example video url```
+    '''
+    
+    formatted_context = "RELEVANT CONTEXT:\n```" + formatted_context + "```"
+    
+    response = client.chat.completions.create(
+        model=model,
+        temperature=temperature,
+        messages=[
+            {"role": "system", "content": instruction},
+            {"role": "user", "content": formatted_context},
+            {"role": "user", "content": question}
+        ]
+    )
+    
+    return response.choices[0].message.content

models/retrieval.py

@@ -0,0 +1,11 @@
+import chromadb
+
+
+def get_relevant_segments(query, db_path, n_results=5):
+    client = chromadb.PersistentClient(db_path)
+    collection= client.get_collection('huberman_videos')
+    
+    results = collection.query(query_texts=[query], n_results=n_results)
+    
+    return results
+

package-lock.json

@@ -0,0 +1,23 @@
+{
+  "name": "ml-qa-chatbot",
+  "lockfileVersion": 3,
+  "requires": true,
+  "packages": {
+    "": {
+      "dependencies": {
+        "dotenv": "^16.4.1"
+      }
+    },
+    "node_modules/dotenv": {
+      "version": "16.4.1",
+      "resolved": "https://registry.npmjs.org/dotenv/-/dotenv-16.4.1.tgz",
+      "integrity": "sha512-CjA3y+Dr3FyFDOAMnxZEGtnW9KBR2M0JvvUtXNW+dYJL5ROWxP9DUHCwgFqpMk0OXCc0ljhaNTr2w/kutYIcHQ==",
+      "engines": {
+        "node": ">=12"
+      },
+      "funding": {
+        "url": "https://github.com/motdotla/dotenv?sponsor=1"
+      }
+    }
+  }
+}

package.json

@@ -0,0 +1,5 @@
+{
+  "dependencies": {
+    "dotenv": "^16.4.1"
+  }
+}

requirements.txt

@@ -0,0 +1,127 @@
+aiofiles==23.2.1
+altair==5.2.0
+annotated-types==0.6.0
+anyio==4.2.0
+asgiref==3.7.2
+attrs==23.2.0
+backoff==2.2.1
+bcrypt==4.1.2
+build==1.0.3
+cachetools==5.3.2
+certifi==2024.2.2
+charset-normalizer==3.3.2
+chroma-hnswlib==0.7.3
+chromadb==0.4.22
+click==8.1.7
+colorama==0.4.6
+coloredlogs==15.0.1
+contourpy==1.2.0
+cycler==0.12.1
+Deprecated==1.2.14
+distro==1.9.0
+fastapi==0.109.2
+ffmpy==0.3.1
+filelock==3.13.1
+flatbuffers==23.5.26
+fonttools==4.47.2
+fsspec==2024.2.0
+google-auth==2.27.0
+googleapis-common-protos==1.62.0
+gradio==4.16.0
+gradio_client==0.8.1
+grpcio==1.60.1
+h11==0.14.0
+httpcore==1.0.2
+httptools==0.6.1
+httpx==0.26.0
+huggingface-hub==0.20.3
+humanfriendly==10.0
+idna==3.6
+importlib-metadata==6.11.0
+importlib-resources==6.1.1
+Jinja2==3.1.3
+joblib==1.3.2
+jsonschema==4.21.1
+jsonschema-specifications==2023.12.1
+kiwisolver==1.4.5
+kubernetes==29.0.0
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+matplotlib==3.8.2
+mdurl==0.1.2
+mmh3==4.1.0
+monotonic==1.6
+mpmath==1.3.0
+networkx==3.2.1
+numpy==1.26.3
+oauthlib==3.2.2
+onnxruntime==1.17.0
+openai==1.11.1
+opentelemetry-api==1.22.0
+opentelemetry-exporter-otlp-proto-common==1.22.0
+opentelemetry-exporter-otlp-proto-grpc==1.22.0
+opentelemetry-instrumentation==0.43b0
+opentelemetry-instrumentation-asgi==0.43b0
+opentelemetry-instrumentation-fastapi==0.43b0
+opentelemetry-proto==1.22.0
+opentelemetry-sdk==1.22.0
+opentelemetry-semantic-conventions==0.43b0
+opentelemetry-util-http==0.43b0
+orjson==3.9.13
+overrides==7.7.0
+packaging==23.2
+pandas==2.2.0
+pillow==10.2.0
+posthog==3.4.0
+protobuf==4.25.2
+pulsar-client==3.4.0
+pyasn1==0.5.1
+pyasn1-modules==0.3.0
+pydantic==2.6.1
+pydantic_core==2.16.2
+pydub==0.25.1
+Pygments==2.17.2
+pyparsing==3.1.1
+PyPika==0.48.9
+pyproject_hooks==1.0.0
+pyreadline3==3.4.1
+python-dateutil==2.8.2
+python-dotenv==1.0.1
+python-multipart==0.0.7
+pytz==2024.1
+PyYAML==6.0.1
+referencing==0.33.0
+regex==2023.12.25
+requests==2.31.0
+requests-oauthlib==1.3.1
+rich==13.7.0
+rpds-py==0.17.1
+rsa==4.9
+ruff==0.2.1
+safetensors==0.4.2
+scikit-learn==1.4.0
+scipy==1.12.0
+semantic-version==2.10.0
+shellingham==1.5.4
+six==1.16.0
+sniffio==1.3.0
+starlette==0.36.3
+sympy==1.12
+tenacity==8.2.3
+threadpoolctl==3.2.0
+tokenizers==0.15.1
+tomlkit==0.12.0
+toolz==0.12.1
+torch==2.2.0
+tqdm==4.66.1
+transformers==4.37.2
+typer==0.9.0
+typing_extensions==4.9.0
+tzdata==2023.4
+urllib3==2.2.0
+uvicorn==0.27.0.post1
+watchfiles==0.21.0
+websocket-client==1.7.0
+websockets==11.0.3
+wrapt==1.16.0
+zipp==3.17.0

run_etl.py

@@ -0,0 +1,14 @@
+from models.etl import run_etl
+
+if __name__ == "__main__":
+    json_path = input("Enter path to JSON file: ")
+    db_path = input("Enter path to database: ")
+    batch_size = int(input("Enter batch size (leave blank for no batching): "))
+
+    if batch_size:
+        overlap = int(input("Enter overlap (leave blank for no overlap): "))
+    else:
+        overlap = None
+
+
+    run_etl(json_path=json_path, db=db_path, batch_size=batch_size, overlap=overlap)

tests/test_main.py

@@ -0,0 +1,43 @@
+from main import run_query
+from utils.general_utils import timeit
+
+@timeit
+def test_with_llm_video():
+    db_path = "data/single_video.db"
+    
+    question = "What are the components of an LLM?"
+    print("Question: ", question)
+    
+    answer = run_query(question,
+                   db_path=db_path
+                   )
+    
+    print("Answer: ", answer)
+
+@timeit
+def test_with_subset():
+    db_path = "data/videos_subset_more_context.db"
+    
+    question = "How should I train for anerobic capacity?"
+    print("Question: ", question)
+    
+    answer = run_query(question,
+                   db_path=db_path,
+                   num_rel_segments=10,
+                   llm_model="gpt-3.5-turbo-0125",
+                   llm_temp=0.1
+                   )
+    
+    print("Answer: ", answer)
+
+
+if __name__ == '__main__':
+    choice = input("Enter 1 for test_with_subset, 2 for test_with_llm_video: ")
+    
+    if choice == "1":
+        test_with_subset()
+    elif choice == "2":
+        test_with_llm_video()
+    else:
+        print("Invalid choice")
+        exit(1)

tests/test_retrieval.py

@@ -0,0 +1,20 @@
+from models.retrieval import get_relevant_segments
+from models.llm import format_context
+from utils.general_utils import timeit
+
+@timeit
+def test_with_subset():
+    db_path = "data/videos_subset.db"
+    
+    question = "What methods can I use to lose weight quickly?"
+    print("Question: ", question)
+    
+    relevant_segments = get_relevant_segments(question,
+                   db_path=db_path
+                   )
+    formatted_segments = format_context(relevant_segments)
+    
+    print("Segments: ", formatted_segments)
+    
+if __name__ == "__main__":
+    test_with_subset()

utils/embedding_utils.py

@@ -0,0 +1,23 @@
+import torch
+from transformers import AutoTokenizer, AutoModel
+from chromadb import Documents, EmbeddingFunction, Embeddings
+
+
+model_name = "YituTech/conv-bert-base"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModel.from_pretrained(model_name)
+
+
+class MyEmbeddingFunction(EmbeddingFunction[Documents]):
+
+    def __call__(self, input: Documents) -> Embeddings:
+        embeddings_list = []
+        
+        for text in input:
+            tokens = tokenizer(text, return_tensors='pt')
+            with torch.no_grad():
+                outputs = model(**tokens)
+            embeddings = outputs.last_hidden_state.mean(dim=1).squeeze().detach().numpy()
+            embeddings_list.append(embeddings)
+        
+        return embeddings_list

utils/general_utils.py

@@ -0,0 +1,12 @@
+import time
+
+def timeit(func):
+    def wrapper(*args, **kwargs):
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        end_time = time.time()
+        execution_time = round(end_time - start_time, 2)
+        print(f"{func.__name__} took {execution_time} seconds.")
+        return result
+    return wrapper
+