added all steps, not fully fuctional

main.ipynb

@@ -2,26 +2,34 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 70,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "from semantic_search import SemanticSearch \n",
     "import pandas as pd\n",
+    "import tiktoken\n",
     "import openai\n",
     "import praw\n",
     "import os\n",
     "import re\n",
     "\n",
-    "searcher = SemanticSearch()\n",
-    "\n",
     "pd.set_option('max_colwidth', 100)\n",
     "pd.set_option('display.max_columns', None)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 71,
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "searcher = SemanticSearch()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -32,16 +40,24 @@
     "reddit = praw.Reddit(client_id=REDDIT_CLIENT_ID, client_secret=REDDIT_CLIENT_SECRET, user_agent=REDDIT_USER_AGENT)"
    ]
   },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Topic Retrieval"
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 72,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "def generate_topics(query, model=\"gpt-3.5-turbo\"):\n",
     "\n",
     "    messages = [\n",
-    "        {\"role\": \"user\", \"content\": f\"Take this query '{query}' and return a list of topics to input in Search so it returns good results. Each topic must stand on its own with respect to the relation of the question.\"},\n",
+    "        {\"role\": \"user\", \"content\": f\"Take this query '{query}' and return a list of short topics to input in Search so it returns good results. Each topic must stand on its own with respect to the relation of the question.\"},\n",
     "    ]\n",
     "\n",
     "    response = openai.ChatCompletion.create(\n",
@@ -58,26 +74,35 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 73,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "['Top-rated pizza places in NYC', 'Best New York-style pizza restaurants', 'Highest reviewed pizzerias in New York City', 'Best pizza joints in Manhattan', 'Award-winning pizza spots in NYC', 'Famous pizza places in New York', 'Top 10 pizza places to try in NYC', 'Authentic pizza restaurants in New York', 'Best pizza slices in NYC', 'Gourmet pizza places in New York City']\n"
-     ]
-    }
-   ],
+   "outputs": [],
+   "source": [
+    "query = \"Where are some nice places where I can work remotely in Malta?\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
    "source": [
-    "query = \"Best pizza place in NYC\"\n",
     "topics = generate_topics(query)\n",
+    "topics = [topic.strip() for topic in topics]\n",
     "print(topics)"
    ]
   },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Relevant Subreddits Retrieval"
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 74,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -88,128 +113,183 @@
     "    topic, limit=200):\n",
     "        posts.append([post.title, post.subreddit, post.selftext])\n",
     "\n",
-    "posts = pd.DataFrame(posts,columns=['title', 'subreddit', 'text'])"
+    "posts = pd.DataFrame(posts,columns=['title', 'subreddit', 'text'])\n",
+    "\n",
+    "# Segments is title, text and subreddit at the end\n",
+    "segments = (posts['title'] + ' ' + posts['subreddit'].astype(str)).tolist()"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 75,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Segments is title, text and subreddit at the end\n",
-    "segments = (posts['title'] + ' ' + posts['subreddit'].astype(str)).tolist()"
+    "searcher.fit(segments, n_neighbors=5)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# TODO: Add distance check here\n",
+    "subreddits = set([result.split()[-1] for result in searcher(query)])\n",
+    "\n",
+    "# Convert to string and \"+\" in between\n",
+    "subreddits = \"+\".join(subreddits)\n",
+    "\n",
+    "print(f\"Relevant subreddits: {subreddits}\")"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Relevant Posts Retrieval"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 76,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Fitting with n=5...\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
-    "searcher.fit(segments)"
+    "segments = []\n",
+    "segment_length = 100\n",
+    "\n",
+    "\n",
+    "for topic in topics:\n",
+    "    for post in reddit.subreddit(subreddits).search(\n",
+    "        topic, limit=50):\n",
+    "            \n",
+    "            comments = \"\"\n",
+    "\n",
+    "            post.comments.replace_more(limit=3)\n",
+    "            for comment in post.comments.list():\n",
+    "                if comment.body != \"[deleted]\":\n",
+    "                    comments += comment.body + \"\\n\"\n",
+    "\n",
+    "            words = comments.split()\n",
+    "            segments.extend([post.title + \" \" + post.id + \"\\n\" + ' '.join(words[i:i+segment_length]) for i in range(0, len(words), segment_length)])"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 77,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "# TODO: Add distance check here\n",
-    "subreddits = set([result.split()[-1] for result in searcher(query)])"
+    "searcher.fit(segments, n_neighbors=5)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Answering the Query"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 78,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Convert to string and \"+\" in between\n",
-    "subreddits = \"+\".join(subreddits)"
+    "def num_tokens(text, model):\n",
+    "    encoding = tiktoken.encoding_for_model(model)\n",
+    "    return len(encoding.encode(text))"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 79,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "final_posts = []\n",
+    "def form_query(query, model, token_budget):\n",
+    "\n",
+    "    relevant_segments = searcher(query)\n",
+    "\n",
+    "    introduction = 'Use the below segments from multiple Reddit posts to answer the subsequent question. If the answer cannot be found in the articles, write \"I could not find an answer.\" Cite each sentence using the [postid] notation found at the start of each segment. Every sentence MUST have a citation!\\n\\n'\n",
+    "\n",
+    "    message = introduction\n",
+    "\n",
+    "    query = f\"\\n\\nQuestion: {query}\"\n",
     "\n",
-    "for post in reddit.subreddit(subreddits).search(\n",
-    "    query, limit=100):\n",
-    "        final_posts.append([post.title, post.subreddit, post.selftext])\n",
+    "    evidence = []\n",
     "\n",
-    "final_posts = pd.DataFrame(final_posts,columns=['title', 'subreddit', 'text'])"
+    "    for i, result in enumerate(relevant_segments):\n",
+    "        if (\n",
+    "            num_tokens(message + result + query, model=model)\n",
+    "            > token_budget\n",
+    "        ):\n",
+    "            break\n",
+    "        else:\n",
+    "            result = result + \"\\n\\n\"\n",
+    "            message += result\n",
+    "            evidence.append(result.split(\"\\n\")[0])\n",
+    "\n",
+    "    evidence = list(set(evidence))\n",
+    "\n",
+    "    return message + query, evidence"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 80,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "final_segments = (final_posts['title'] + ' ' + final_posts['subreddit'].astype(str)).tolist()"
+    "def generate_answer(query, model, token_budget, temperature):\n",
+    "    \n",
+    "    message, evidence = form_query(query, model, token_budget)\n",
+    "\n",
+    "    messages = [\n",
+    "        {\"role\": \"user\", \"content\": message},\n",
+    "    ]\n",
+    "\n",
+    "    print(message)\n",
+    "\n",
+    "    response = openai.ChatCompletion.create(\n",
+    "        model=model,\n",
+    "        messages=messages,\n",
+    "        temperature=temperature\n",
+    "    )\n",
+    "    \n",
+    "    response_message = response[\"choices\"][0][\"message\"][\"content\"]\n",
+    "\n",
+    "    return response_message, evidence"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 81,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "final_searcher = SemanticSearch()"
+    "answer, evidence = generate_answer(query, \"gpt-3.5-turbo\", 1000, 0)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 82,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Fitting with n=5...\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
-    "final_searcher.fit(final_segments)"
+    "query"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 83,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "['Best Pizza in Manhattan? FoodNYC',\n",
-       " 'Best Pizza slice in NYC? FoodNYC',\n",
-       " 'Best Pizza, Manhattan NYC at 40.713050, -74.007230 The8BitRyanReddit',\n",
-       " \"Controversial: What's the best pizza place in all of NYC? newyorkcity\",\n",
-       " \"Best pizza place in NYC that tourists DON'T know? circlejerknyc\"]"
-      ]
-     },
-     "execution_count": 83,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
-    "final_searcher(query)"
+    "answer"
    ]
   }
  ],

semantic_search.py

@@ -21,12 +21,12 @@ class SemanticSearch:
     
     def __call__(self, text, return_data=True):
         inp_emb = self.use([text])
-        neighbors = self.nn.kneighbors(inp_emb, return_distance=False)[0]
+        neighbors, distances = self.nn.kneighbors(inp_emb, return_distance=True)[0]
         
         if return_data:
-            return [self.data[i] for i in neighbors]
+            return [self.data[i] for i in neighbors], distances
         else:
-            return neighbors
+            return neighbors, distances
     
     
     def get_text_embedding(self, texts, batch=1000):