method works well

main.ipynb

@@ -2,13 +2,18 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 30,
+   "execution_count": 70,
    "metadata": {},
    "outputs": [],
    "source": [
+    "from semantic_search import SemanticSearch \n",
     "import pandas as pd\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)"
@@ -16,7 +21,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 71,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -29,19 +34,182 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 233,
+   "execution_count": 72,
+   "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",
+    "    ]\n",
+    "\n",
+    "    response = openai.ChatCompletion.create(\n",
+    "        model=model,\n",
+    "        messages=messages\n",
+    "    )\n",
+    "\n",
+    "    response_message = response[\"choices\"][0][\"message\"][\"content\"]\n",
+    "\n",
+    "    topics = re.sub(r'^\\d+\\.\\s*', '', response_message, flags=re.MULTILINE).split(\"\\n\")\n",
+    "\n",
+    "    return topics"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 73,
+   "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"
+     ]
+    }
+   ],
+   "source": [
+    "query = \"Best pizza place in NYC\"\n",
+    "topics = generate_topics(query)\n",
+    "print(topics)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 74,
    "metadata": {},
    "outputs": [],
    "source": [
     "posts = []\n",
-    "comments = []\n",
     "\n",
-    "query = \"Which place sells the best pastizzis in Malta?\"\n",
+    "for topic in topics:\n",
+    "    for post in reddit.subreddit(\"all\").search(\n",
+    "    topic, limit=200):\n",
+    "        posts.append([post.title, post.subreddit, post.selftext])\n",
+    "\n",
+    "posts = pd.DataFrame(posts,columns=['title', 'subreddit', 'text'])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 75,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Segments is title, text and subreddit at the end\n",
+    "segments = (posts['title'] + ' ' + posts['subreddit'].astype(str)).tolist()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 76,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Fitting with n=5...\n"
+     ]
+    }
+   ],
+   "source": [
+    "searcher.fit(segments)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 77,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# TODO: Add distance check here\n",
+    "subreddits = set([result.split()[-1] for result in searcher(query)])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 78,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Convert to string and \"+\" in between\n",
+    "subreddits = \"+\".join(subreddits)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 79,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "final_posts = []\n",
     "\n",
-    "for post in reddit.subreddit(\"all\").search(\"Malta cuisine\", limit=200):\n",
-    "    posts.append([post.title, post.subreddit, post.selftext, post.created])\n",
+    "for post in reddit.subreddit(subreddits).search(\n",
+    "    query, limit=100):\n",
+    "        final_posts.append([post.title, post.subreddit, post.selftext])\n",
     "\n",
-    "posts = pd.DataFrame(posts,columns=['title', 'subreddit', 'text', 'created'])"
+    "final_posts = pd.DataFrame(final_posts,columns=['title', 'subreddit', 'text'])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 80,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "final_segments = (final_posts['title'] + ' ' + final_posts['subreddit'].astype(str)).tolist()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 81,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "final_searcher = SemanticSearch()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 82,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Fitting with n=5...\n"
+     ]
+    }
+   ],
+   "source": [
+    "final_searcher.fit(final_segments)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 83,
+   "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"
+    }
+   ],
+   "source": [
+    "final_searcher(query)"
    ]
   }
  ],

semantic_search.py

@@ -0,0 +1,39 @@
+from sklearn.neighbors import NearestNeighbors
+import tensorflow_hub as hub
+import numpy as np
+
+class SemanticSearch:
+    
+    def __init__(self):
+        self.use = hub.load('https://tfhub.dev/google/universal-sentence-encoder/4')
+        self.fitted = False
+    
+    
+    def fit(self, data, batch=1000, n_neighbors=5):
+        print(f"Fitting with n={n_neighbors}...")
+        self.data = data
+        self.embeddings = self.get_text_embedding(data, batch=batch)
+        n_neighbors = min(n_neighbors, len(self.embeddings))
+        self.nn = NearestNeighbors(n_neighbors=n_neighbors)
+        self.nn.fit(self.embeddings)
+        self.fitted = True
+    
+    
+    def __call__(self, text, return_data=True):
+        inp_emb = self.use([text])
+        neighbors = self.nn.kneighbors(inp_emb, return_distance=False)[0]
+        
+        if return_data:
+            return [self.data[i] for i in neighbors]
+        else:
+            return neighbors
+    
+    
+    def get_text_embedding(self, texts, batch=1000):
+        embeddings = []
+        for i in range(0, len(texts), batch):
+            text_batch = texts[i:(i+batch)]
+            emb_batch = self.use(text_batch)
+            embeddings.append(emb_batch)
+        embeddings = np.vstack(embeddings)
+        return embeddings