File size: 37,180 Bytes
4aa3246 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 |
{
"cells": [
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/Users/ladka6/Projects/semantic-search/venv/lib/python3.9/site-packages/urllib3/__init__.py:34: NotOpenSSLWarning: urllib3 v2 only supports OpenSSL 1.1.1+, currently the 'ssl' module is compiled with 'LibreSSL 2.8.3'. See: https://github.com/urllib3/urllib3/issues/3020\n",
" warnings.warn(\n"
]
}
],
"source": [
"import re\n",
"import nltk\n",
"import numpy as np\n",
"import pandas as pd\n",
"from langdetect import detect\n",
"from sentence_transformers import SentenceTransformer, InputExample, losses\n",
"from sentence_transformers.util import semantic_search \n",
"from torch.utils.data import DataLoader\n",
"from nltk.corpus import stopwords\n",
"from nltk.tokenize import word_tokenize\n",
"from nltk.stem import PorterStemmer\n",
"from sklearn.preprocessing import MultiLabelBinarizer\n",
"import faiss\n",
"from FlagEmbedding import FlagReranker\n",
"from sklearn.model_selection import train_test_split\n"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"FILE_PATH = '../data/book_data.csv'\n",
"\n",
"df = pd.read_csv(FILE_PATH)\n",
"\n",
"df = df.dropna(subset=['Title', 'Description', 'Genres'])\n"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"def preprocess_text(text):\n",
" text = text.lower()\n",
" text = re.sub(r'[^a-zA-Z\\s]', '', text)\n",
" tokens = word_tokenize(text)\n",
" stop_words = set(stopwords.words('english'))\n",
" tokens = [token for token in tokens if token not in stop_words]\n",
" text = ' '.join(tokens)\n",
" return text\n"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"# def title_description_genres_to_string(row):\n",
"# genres = ' '.join([genre.strip() for genre in row['Genres'].split(',')])\n",
"\n",
"# title = row['Title']\n",
"\n",
"# descriptions = str(row['Description'])\n",
"\n",
"# tokens = word_tokenize(descriptions)\n",
"\n",
"# tokens = [word.lower() for word in tokens if word.isalpha()]\n",
"\n",
"# stop_words = set(stopwords.words('english'))\n",
"# tokens = [word for word in tokens if word not in stop_words]\n",
"\n",
"# porter = PorterStemmer()\n",
"# tokens = [porter.stem(word) for word in tokens]\n",
"\n",
"# preprocessed_text = ' '.join(tokens)\n",
"\n",
"# return \"%s %s %s\" %(title, genres, preprocessed_text)\n"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"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>Id</th>\n",
" <th>Title</th>\n",
" <th>Author</th>\n",
" <th>Rating</th>\n",
" <th>Description</th>\n",
" <th>Genres</th>\n",
" <th>Reviews</th>\n",
" <th>Combined</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>0</th>\n",
" <td>1</td>\n",
" <td>Beowulf</td>\n",
" <td>Seamus Heaney</td>\n",
" <td>3.48</td>\n",
" <td>Composed toward the end of the first millenniu...</td>\n",
" <td>Classics, Poetry, Fiction, Fantasy, Mythology,...</td>\n",
" <td>*bum bum* IN A WORLD . . . *bum bum* . . . FUL...</td>\n",
" <td>Beowulf Classics Poetry Fiction Fantasy Mythol...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>1</th>\n",
" <td>2</td>\n",
" <td>The Evening and the Morning</td>\n",
" <td>Ken Follett</td>\n",
" <td>4.38</td>\n",
" <td>The thrilling and addictive prequel to The Pil...</td>\n",
" <td>Historical Fiction, Fiction, Historical, Audio...</td>\n",
" <td>It's 997 CE, the end of the Dark Ages in Engla...</td>\n",
" <td>The Evening and the Morning Historical Fiction...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2</th>\n",
" <td>3</td>\n",
" <td>The Abbot's Tale</td>\n",
" <td>Conn Iggulden</td>\n",
" <td>4.05</td>\n",
" <td>In the year 937, the new king of England, a gr...</td>\n",
" <td>Historical Fiction, Fiction, Historical, Medie...</td>\n",
" <td>There is never one truth, one love, or one ene...</td>\n",
" <td>The Abbot's Tale Historical Fiction Fiction Hi...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>3</th>\n",
" <td>4</td>\n",
" <td>Ibn Fadlān and the Land of Darkness: Arab Trav...</td>\n",
" <td>Ahmad ibn Fadlān</td>\n",
" <td>3.87</td>\n",
" <td>In 922 AD, an Arab envoy from Baghdad named Ib...</td>\n",
" <td>History, Travel, Nonfiction, Classics, Islam, ...</td>\n",
" <td>رسالة ابن فضلان .. أو ما يسمى برحلة ابن فضلان ...</td>\n",
" <td>Ibn Fadlān and the Land of Darkness: Arab Trav...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>4</th>\n",
" <td>5</td>\n",
" <td>The Empty Throne</td>\n",
" <td>Bernard Cornwell</td>\n",
" <td>4.38</td>\n",
" <td>This eighth entry in New York Times bestsellin...</td>\n",
" <td>Historical Fiction, Fiction, Historical, Medie...</td>\n",
" <td>The Empty Throne was an improvement over The P...</td>\n",
" <td>The Empty Throne Historical Fiction Fiction Hi...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>...</th>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" <td>...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>1780</th>\n",
" <td>1843</td>\n",
" <td>The Soul of an Octopus</td>\n",
" <td>Sy Montgomery</td>\n",
" <td>3.93</td>\n",
" <td>In pursuit of the wild, solitary, predatory oc...</td>\n",
" <td>Nonfiction, Science, Animals, Nature, Memoir, ...</td>\n",
" <td>I'm kind of \"eh\" on this book. It bills itself...</td>\n",
" <td>The Soul of an Octopus Nonfiction Science Anim...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>1781</th>\n",
" <td>1844</td>\n",
" <td>The Jim Corbett Omnibus: \"Man-eaters of Kumaon...</td>\n",
" <td>Jim Corbett</td>\n",
" <td>4.54</td>\n",
" <td>Jim Corbett's riveting accounts of shikar in t...</td>\n",
" <td>Nonfiction, Wildlife, Nature, Travel, Biograph...</td>\n",
" <td>One of the best books ever written, this book ...</td>\n",
" <td>The Jim Corbett Omnibus: \"Man-eaters of Kumaon...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>1782</th>\n",
" <td>1845</td>\n",
" <td>The Tiger: A True Story of Vengeance and Survival</td>\n",
" <td>John Vaillant</td>\n",
" <td>4.07</td>\n",
" <td>It’s December 1997, and a man-eating tiger is ...</td>\n",
" <td>Nonfiction, History, Nature, Animals, Russia, ...</td>\n",
" <td>Fearful symmetry indeed. In 1997, during time ...</td>\n",
" <td>The Tiger: A True Story of Vengeance and Survi...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>1783</th>\n",
" <td>1846</td>\n",
" <td>100 Heartbeats: The Race to Save Earth's Most ...</td>\n",
" <td>Jeff Corwin</td>\n",
" <td>4.17</td>\n",
" <td>It's no secret that our planet―and the delicat...</td>\n",
" <td>Nonfiction, Animals, Science, Nature, Conserva...</td>\n",
" <td>I learned a lot reading this book. Frightening...</td>\n",
" <td>100 Heartbeats: The Race to Save Earth's Most ...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>1784</th>\n",
" <td>1847</td>\n",
" <td>An Elephant in My Kitchen</td>\n",
" <td>Françoise Malby-Anthony</td>\n",
" <td>4.39</td>\n",
" <td>A blonde, chic Parisienne, Francoise never exp...</td>\n",
" <td>Nonfiction, Animals, Memoir, Africa, Nature, B...</td>\n",
" <td>Thula Thula, South Africa, the sanctuary for e...</td>\n",
" <td>An Elephant in My Kitchen Nonfiction Animals M...</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"<p>1779 rows × 8 columns</p>\n",
"</div>"
],
"text/plain": [
" Id Title \\\n",
"0 1 Beowulf \n",
"1 2 The Evening and the Morning \n",
"2 3 The Abbot's Tale \n",
"3 4 Ibn Fadlān and the Land of Darkness: Arab Trav... \n",
"4 5 The Empty Throne \n",
"... ... ... \n",
"1780 1843 The Soul of an Octopus \n",
"1781 1844 The Jim Corbett Omnibus: \"Man-eaters of Kumaon... \n",
"1782 1845 The Tiger: A True Story of Vengeance and Survival \n",
"1783 1846 100 Heartbeats: The Race to Save Earth's Most ... \n",
"1784 1847 An Elephant in My Kitchen \n",
"\n",
" Author Rating \\\n",
"0 Seamus Heaney 3.48 \n",
"1 Ken Follett 4.38 \n",
"2 Conn Iggulden 4.05 \n",
"3 Ahmad ibn Fadlān 3.87 \n",
"4 Bernard Cornwell 4.38 \n",
"... ... ... \n",
"1780 Sy Montgomery 3.93 \n",
"1781 Jim Corbett 4.54 \n",
"1782 John Vaillant 4.07 \n",
"1783 Jeff Corwin 4.17 \n",
"1784 Françoise Malby-Anthony 4.39 \n",
"\n",
" Description \\\n",
"0 Composed toward the end of the first millenniu... \n",
"1 The thrilling and addictive prequel to The Pil... \n",
"2 In the year 937, the new king of England, a gr... \n",
"3 In 922 AD, an Arab envoy from Baghdad named Ib... \n",
"4 This eighth entry in New York Times bestsellin... \n",
"... ... \n",
"1780 In pursuit of the wild, solitary, predatory oc... \n",
"1781 Jim Corbett's riveting accounts of shikar in t... \n",
"1782 It’s December 1997, and a man-eating tiger is ... \n",
"1783 It's no secret that our planet―and the delicat... \n",
"1784 A blonde, chic Parisienne, Francoise never exp... \n",
"\n",
" Genres \\\n",
"0 Classics, Poetry, Fiction, Fantasy, Mythology,... \n",
"1 Historical Fiction, Fiction, Historical, Audio... \n",
"2 Historical Fiction, Fiction, Historical, Medie... \n",
"3 History, Travel, Nonfiction, Classics, Islam, ... \n",
"4 Historical Fiction, Fiction, Historical, Medie... \n",
"... ... \n",
"1780 Nonfiction, Science, Animals, Nature, Memoir, ... \n",
"1781 Nonfiction, Wildlife, Nature, Travel, Biograph... \n",
"1782 Nonfiction, History, Nature, Animals, Russia, ... \n",
"1783 Nonfiction, Animals, Science, Nature, Conserva... \n",
"1784 Nonfiction, Animals, Memoir, Africa, Nature, B... \n",
"\n",
" Reviews \\\n",
"0 *bum bum* IN A WORLD . . . *bum bum* . . . FUL... \n",
"1 It's 997 CE, the end of the Dark Ages in Engla... \n",
"2 There is never one truth, one love, or one ene... \n",
"3 رسالة ابن فضلان .. أو ما يسمى برحلة ابن فضلان ... \n",
"4 The Empty Throne was an improvement over The P... \n",
"... ... \n",
"1780 I'm kind of \"eh\" on this book. It bills itself... \n",
"1781 One of the best books ever written, this book ... \n",
"1782 Fearful symmetry indeed. In 1997, during time ... \n",
"1783 I learned a lot reading this book. Frightening... \n",
"1784 Thula Thula, South Africa, the sanctuary for e... \n",
"\n",
" Combined \n",
"0 Beowulf Classics Poetry Fiction Fantasy Mythol... \n",
"1 The Evening and the Morning Historical Fiction... \n",
"2 The Abbot's Tale Historical Fiction Fiction Hi... \n",
"3 Ibn Fadlān and the Land of Darkness: Arab Trav... \n",
"4 The Empty Throne Historical Fiction Fiction Hi... \n",
"... ... \n",
"1780 The Soul of an Octopus Nonfiction Science Anim... \n",
"1781 The Jim Corbett Omnibus: \"Man-eaters of Kumaon... \n",
"1782 The Tiger: A True Story of Vengeance and Survi... \n",
"1783 100 Heartbeats: The Race to Save Earth's Most ... \n",
"1784 An Elephant in My Kitchen Nonfiction Animals M... \n",
"\n",
"[1779 rows x 8 columns]"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# df['Combined'] = df.apply(title_description_genres_to_string, axis=1)\n",
"df['Title'] = preprocess_text(df['Title'])\n",
"df['Description'] = preprocess_text(df['Description'])\n",
"df['Genres'] = preprocess_text(df['Genres'])\n"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "6295f3bbee6d43c1a32f8e626eaa31cc",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch: 0%| | 0/1 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "2c6f9ff852e24f5fb87826524a452f65",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Iteration: 0%| | 0/45 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"train, valid = train_test_split(df, test_size=0.2, random_state=42)\n",
"\n",
"train_examples = [\n",
" InputExample(texts=[row[\"Title\"], row['Genres'], row['Description']], label=1.0) for _, row in train.iterrows()\n",
"]\n",
"\n",
"model = SentenceTransformer(\"paraphrase-MiniLM-L6-v2\")\n",
"\n",
"data_loader = DataLoader(train_examples, shuffle=True, batch_size=32)\n",
"train_loss = losses.CosineSimilarityLoss(model=model)\n",
"\n",
"model.fit(train_objectives=[(data_loader, train_loss)], epochs=1)\n",
"\n",
"model.save(\"out/fine_tuned_sbert_model_test\")\n"
]
},
{
"cell_type": "code",
"execution_count": 48,
"metadata": {},
"outputs": [],
"source": [
"documents = list(df['Description'])\n",
"\n",
"sbert_model = SentenceTransformer('out/fine_tuned_sbert_model_test')\n",
"\n",
"document_embeddings = sbert_model.encode(documents, convert_to_tensor=True)\n",
"\n",
"index = faiss.IndexFlatL2(document_embeddings.size(1))\n",
"\n",
"index.add(document_embeddings.cpu().numpy())\n",
"\n",
"faiss.write_index(index, 'vectors/fine_tuned_faiss_index_t.index')\n"
]
},
{
"cell_type": "code",
"execution_count": 53,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"{0: array([1192, 1188, 218, 280, 645]), 1: array([426, 394, 428, 422, 393]), 2: array([1474, 59, 498, 336, 1376]), 3: array([765, 798, 790, 170, 809]), 4: array([ 485, 664, 1311, 1035, 1328]), 5: array([ 479, 729, 466, 1521, 828]), 6: array([ 109, 114, 101, 145, 1660]), 7: array([ 774, 767, 165, 1647, 305]), 8: array([ 85, 1431, 907, 879, 910]), 9: array([1188, 1035, 1019, 1114, 1192]), 10: array([1721, 674, 1303, 685, 709]), 11: array([530, 550, 527, 515, 554]), 12: array([984, 975, 450, 452, 449]), 13: array([ 538, 1367, 1645, 534, 59]), 14: array([1057, 1019, 1353, 1039, 1015]), 15: array([ 677, 1319, 664, 1311, 688]), 16: array([1101, 1625, 1068, 764, 664]), 17: array([316, 303, 294, 683, 390]), 18: array([1149, 1157, 862, 711, 351]), 19: array([ 930, 1035, 1019, 1237, 266]), 20: array([ 344, 313, 1285, 1185, 1767]), 21: array([ 590, 391, 1485, 196, 48]), 22: array([948, 952, 916, 943, 755]), 23: array([785, 798, 802, 769, 809]), 24: array([ 881, 341, 1651, 816, 1369]), 25: array([ 561, 516, 532, 1378, 1763]), 26: array([1140, 1114, 1138, 134, 1506]), 27: array([303, 316, 294, 683, 390]), 28: array([ 220, 844, 1271, 1520, 1130]), 29: array([1271, 1131, 1266, 1200, 1154]), 30: array([1080, 1085, 1067, 1103, 1089]), 31: array([ 741, 728, 274, 717, 1500]), 32: array([1268, 1277, 1261, 1259, 1264]), 33: array([599, 996, 232, 970, 571]), 34: array([324, 311, 294, 364, 59]), 35: array([1711, 1699, 692, 371, 629]), 36: array([942, 926, 932, 943, 958]), 37: array([ 548, 545, 533, 515, 1374]), 38: array([ 829, 815, 1454, 834, 842]), 39: array([1366, 1378, 1360, 1368, 1406]), 40: array([ 381, 656, 716, 1651, 1047]), 41: array([1450, 1449, 1447, 1433, 1428]), 42: array([ 29, 842, 1061, 846, 666]), 43: array([865, 897, 968, 957, 891]), 44: array([1567, 1560, 1583, 684, 708]), 45: array([1043, 1047, 1051, 1163, 1198]), 46: array([1492, 1519, 1505, 987, 262]), 47: array([1709, 1707, 1708, 1287, 1278]), 48: array([1770, 101, 1767, 137, 470]), 49: array([1716, 1697, 1718, 1720, 1733]), 50: array([1253, 1234, 1246, 1229, 1231]), 51: array([115, 118, 307, 74, 126]), 52: array([ 679, 1149, 652, 602, 1035]), 53: array([1423, 1425, 1415, 839, 1051]), 54: array([1244, 1229, 1254, 1237, 701]), 55: array([1090, 1099, 1069, 1071, 1096]), 56: array([ 65, 482, 1434, 1614, 101]), 57: array([ 212, 232, 493, 869, 1451]), 58: array([ 557, 1681, 906, 1646, 1647]), 59: array([1110, 1103, 1067, 1100, 1088]), 60: array([1616, 1602, 1598, 1617, 1603]), 61: array([1252, 1218, 1034, 1235, 1461]), 62: array([ 70, 497, 1288, 307, 343]), 63: array([ 339, 343, 1500, 1521, 101]), 64: array([366, 391, 124, 369, 368]), 65: array([1664, 1439, 1651, 1663, 1368]), 66: array([ 483, 463, 505, 1349, 664]), 67: array([ 576, 819, 827, 836, 1281]), 68: array([135, 246, 59, 498, 270]), 69: array([23, 17, 46, 38, 16]), 70: array([ 628, 652, 1035, 632, 1133]), 71: array([ 886, 666, 676, 688, 1352]), 72: array([1237, 1240, 1250, 1227, 672]), 73: array([ 354, 89, 1328, 1017, 392]), 74: array([1033, 1032, 1013, 1039, 1263]), 75: array([1344, 1330, 1336, 1333, 1345]), 76: array([1512, 1498, 1511, 1505, 1508]), 77: array([ 352, 469, 1013, 610, 624]), 78: array([ 631, 1411, 1413, 1183, 1181]), 79: array([1636, 1653, 1648, 1640, 186]), 80: array([ 478, 856, 1521, 127, 1263]), 81: array([ 383, 124, 1126, 49, 1031]), 82: array([1030, 1281, 1733, 1040, 1732]), 83: array([237, 223, 206, 198, 226]), 84: array([1358, 1363, 1378, 1381, 1602]), 85: array([ 331, 1273, 88, 106, 688]), 86: array([1455, 177, 12, 43, 1406]), 87: array([111, 369, 567, 720, 725]), 88: array([1720, 1697, 1716, 1719, 1718]), 89: array([1091, 1097, 1092, 1081, 1089]), 90: array([188, 194, 182, 147, 156]), 91: array([ 651, 1167, 1021, 1168, 844]), 92: array([936, 926, 916, 942, 913]), 93: array([413, 433, 441, 446, 396]), 94: array([464, 465, 492, 505, 474]), 95: array([ 598, 712, 600, 683, 1018]), 96: array([ 939, 1491, 744, 279, 945]), 97: array([367, 366, 738, 213, 218]), 98: array([ 162, 1089, 761, 12, 1390]), 99: array([ 332, 268, 1044, 1154, 1562]), 100: array([1641, 59, 498, 1462, 320]), 101: array([1755, 1703, 1263, 1274, 1723]), 102: array([1679, 450, 984, 966, 1002]), 103: array([1068, 1098, 1097, 1092, 1103]), 104: array([173, 194, 174, 165, 172]), 105: array([1151, 617, 1138, 1125, 392]), 106: array([453, 967, 452, 449, 451]), 107: array([ 754, 758, 723, 274, 1491]), 108: array([1698, 1717, 239, 1241, 236]), 109: array([1026, 1035, 319, 656, 1651]), 110: array([1205, 1183, 1180, 1198, 631]), 111: array([203, 746, 196, 720, 216]), 112: array([ 862, 1149, 1295, 711, 821]), 113: array([ 374, 842, 703, 1732, 1017]), 114: array([1440, 1439, 1362, 1651, 609]), 115: array([471, 492, 467, 505, 463]), 116: array([1336, 1333, 1317, 1344, 1330]), 117: array([486, 504, 17, 492, 38]), 118: array([1269, 1701, 651, 1198, 1197]), 119: array([ 585, 568, 118, 26, 1651]), 120: array([582, 913, 80, 656, 89]), 121: array([425, 442, 447, 395, 422]), 122: array([1276, 1048, 1035, 1396, 1294]), 123: array([1393, 1359, 1363, 1378, 1369]), 124: array([1712, 1713, 725, 74, 1288]), 125: array([1586, 1585, 1611, 1626, 1606]), 126: array([ 851, 1441, 829, 43, 838]), 127: array([1407, 1378, 1364, 1363, 1602]), 128: array([1277, 1268, 1297, 1263, 1264]), 129: array([1023, 629, 1034, 1017, 381]), 130: array([1368, 1406, 1362, 1363, 550]), 131: array([ 49, 383, 363, 59, 498]), 132: array([ 994, 970, 461, 12, 1500]), 133: array([1442, 764, 672, 1285, 676]), 134: array([1381, 1363, 1370, 533, 1358]), 135: array([618, 624, 662, 655, 619]), 136: array([ 534, 1367, 521, 1760, 549]), 137: array([ 710, 1287, 1035, 1295, 567]), 138: array([1461, 1252, 308, 593, 1059]), 139: array([1744, 334, 1746, 64, 1772]), 140: array([ 867, 1448, 1447, 1449, 1450]), 141: array([ 78, 1432, 15, 26, 242]), 142: array([ 210, 1485, 227, 1462, 1027]), 143: array([1208, 632, 664, 1311, 656]), 144: array([ 73, 838, 593, 829, 816]), 145: array([ 69, 335, 510, 1583, 186]), 146: array([ 576, 819, 827, 836, 1281]), 147: array([1181, 1182, 1184, 1167, 1051]), 148: array([1114, 1651, 1188, 1521, 645]), 149: array([730, 728, 753, 724, 717]), 150: array([ 526, 558, 113, 1652, 1472]), 151: array([433, 422, 439, 444, 410]), 152: array([ 450, 975, 984, 966, 1002]), 153: array([ 607, 567, 1729, 391, 838]), 154: array([1384, 1722, 729, 1763, 1407]), 155: array([1494, 719, 1509, 1208, 762]), 156: array([ 76, 301, 137, 470, 43]), 157: array([1486, 319, 217, 242, 364]), 158: array([1730, 1696, 1278, 1274, 1723]), 159: array([1316, 1354, 252, 1313, 1017]), 160: array([ 65, 482, 1434, 1614, 101]), 161: array([ 514, 1768, 527, 525, 1368]), 162: array([ 298, 1275, 1732, 1121, 346]), 163: array([ 610, 1263, 81, 838, 923]), 164: array([ 96, 254, 228, 1268, 1277]), 165: array([1000, 983, 970, 974, 1008]), 166: array([978, 725, 221, 196, 607]), 167: array([ 30, 1471, 1007, 1464, 1256]), 168: array([1128, 1146, 1147, 382, 636]), 169: array([ 669, 1267, 1295, 1294, 739]), 170: array([1472, 1650, 1676, 1647, 134]), 171: array([ 351, 346, 1732, 1149, 1017]), 172: array([1428, 1450, 1449, 1433, 1446]), 173: array([1376, 972, 1179, 59, 498]), 174: array([ 99, 1427, 1623, 218, 280]), 175: array([ 914, 134, 101, 506, 1521]), 176: array([247, 120, 245, 269, 487]), 177: array([1380, 1653, 1472, 1407, 664]), 178: array([ 834, 829, 1297, 43, 1275]), 179: array([ 140, 251, 1651, 300, 664]), 180: array([ 168, 147, 182, 155, 1359]), 181: array([923, 921, 319, 567, 603]), 182: array([ 239, 236, 209, 240, 1679]), 183: array([1111, 1542, 433, 400, 459]), 184: array([ 694, 1275, 1396, 708, 1045]), 185: array([1706, 1279, 1687, 1291, 1691]), 186: array([ 462, 1481, 1003, 87, 1466]), 187: array([ 820, 1453, 847, 43, 829]), 188: array([ 532, 514, 1768, 525, 1763]), 189: array([1124, 1119, 1141, 1112, 1159]), 190: array([32, 11, 2, 46, 7]), 191: array([1345, 1314, 1330, 1344, 1333]), 192: array([233, 211, 206, 217, 201]), 193: array([ 70, 497, 1288, 307, 343]), 194: array([1121, 1732, 1149, 298, 839]), 195: array([ 909, 1359, 1387, 1360, 667]), 196: array([932, 942, 916, 943, 926]), 197: array([185, 165, 148, 173, 151]), 198: array([1701, 1303, 1284, 1035, 835]), 199: array([1377, 1363, 1372, 1406, 1360]), 200: array([ 342, 1521, 493, 1523, 134]), 201: array([ 44, 755, 1506, 347, 567]), 202: array([ 900, 1368, 514, 1768, 1362]), 203: array([ 570, 602, 569, 1235, 196]), 204: array([ 973, 204, 1485, 196, 599]), 205: array([ 889, 815, 1454, 347, 957]), 206: array([ 196, 720, 377, 723, 1485]), 207: array([ 879, 829, 672, 1353, 1732]), 208: array([ 987, 970, 982, 1508, 1006]), 209: array([1342, 1314, 1333, 1344, 1345]), 210: array([ 59, 498, 364, 1035, 270]), 211: array([ 58, 1209, 1211, 1251, 672]), 212: array([289, 256, 300, 261, 270]), 213: array([ 297, 1240, 307, 333, 319]), 214: array([1113, 1149, 1129, 1271, 1130]), 215: array([1073, 1082, 1098, 1071, 1065]), 216: array([1565, 1573, 1555, 1680, 582]), 217: array([1320, 1035, 638, 1309, 1240]), 218: array([1433, 1449, 1450, 1447, 1446]), 219: array([1561, 571, 858, 810, 859]), 220: array([ 123, 101, 342, 1521, 134]), 221: array([1353, 693, 706, 689, 677]), 222: array([ 555, 672, 514, 1768, 1760]), 223: array([1639, 1648, 1653, 1636, 1651]), 224: array([270, 134, 59, 498, 664]), 225: array([1460, 1466, 1472, 1478, 113]), 226: array([1133, 652, 1035, 1157, 677]), 227: array([1615, 1594, 1626, 761, 1162]), 228: array([1599, 1617, 1590, 1603, 264]), 229: array([813, 818, 822, 43, 829]), 230: array([1580, 761, 1554, 1060, 1543]), 231: array([ 275, 318, 1524, 256, 892]), 232: array([ 815, 1454, 829, 43, 1285]), 233: array([963, 969, 965, 983, 970]), 234: array([350, 322, 308, 593, 837]), 235: array([772, 809, 802, 790, 797]), 236: array([1541, 1339, 1097, 664, 1311]), 237: array([1527, 218, 280, 127, 1100]), 238: array([ 964, 1000, 971, 970, 983]), 239: array([ 999, 1160, 1157, 1155, 1125]), 240: array([1281, 1264, 1294, 1278, 1295]), 241: array([1278, 1287, 1298, 1297, 1264]), 242: array([1696, 1730, 1694, 1722, 1278]), 243: array([ 65, 482, 1434, 1614, 101]), 244: array([1745, 1407, 1754, 1763, 560]), 245: array([ 67, 1445, 59, 498, 375]), 246: array([1524, 1521, 249, 1288, 1520]), 247: array([1107, 1067, 849, 461, 1094]), 248: array([1548, 1573, 1541, 1543, 1339]), 249: array([ 124, 364, 1017, 59, 498]), 250: array([543, 549, 516, 305, 80]), 251: array([1041, 1699, 1132, 1268, 1017]), 252: array([ 78, 1432, 15, 26, 242]), 253: array([1343, 1335, 688, 1352, 1212]), 254: array([1530, 363, 723, 196, 1525]), 255: array([1223, 1222, 1226, 1232, 1240]), 256: array([ 873, 274, 1027, 1499, 1015]), 257: array([1653, 1636, 1648, 1640, 1639]), 258: array([1741, 895, 1521, 1760, 1740]), 259: array([ 591, 720, 347, 1521, 725]), 260: array([1181, 1182, 1184, 1167, 1051]), 261: array([1714, 1284, 1303, 1733, 1297]), 262: array([1360, 1363, 533, 1372, 1378]), 263: array([1405, 1256, 1243, 1385, 1404]), 264: array([846, 856, 839, 29, 838]), 265: array([420, 402, 394, 554, 395]), 266: array([ 637, 59, 498, 1188, 1017]), 267: array([1488, 1516, 57, 1499, 213]), 268: array([ 707, 1346, 693, 677, 1319]), 269: array([ 527, 514, 1768, 550, 545]), 270: array([1288, 844, 1521, 725, 1035]), 271: array([ 261, 1651, 260, 664, 1311]), 272: array([777, 802, 793, 797, 798]), 273: array([ 240, 209, 236, 239, 1679]), 274: array([1169, 216, 755, 1060, 493]), 275: array([ 952, 1185, 916, 936, 942]), 276: array([ 925, 1019, 942, 1264, 1175]), 277: array([ 820, 1453, 847, 43, 829]), 278: array([593, 838, 920, 73, 953]), 279: array([415, 423, 432, 412, 435]), 280: array([ 71, 911, 101, 1521, 12]), 281: array([365, 379, 363, 672, 701]), 282: array([ 259, 1521, 1651, 1767, 101]), 283: array([ 51, 512, 1034, 313, 215]), 284: array([1225, 1220, 1351, 1223, 1222]), 285: array([ 910, 957, 908, 1359, 525]), 286: array([1659, 1645, 1435, 1642, 1126]), 287: array([1204, 1206, 278, 1195, 1203]), 288: array([916, 943, 932, 944, 942]), 289: array([ 371, 86, 1294, 319, 664]), 290: array([1704, 1303, 1284, 1708, 1297]), 291: array([432, 394, 414, 441, 446]), 292: array([1255, 664, 1311, 493, 1237]), 293: array([408, 407, 435, 437, 423]), 294: array([1364, 1407, 1363, 1396, 1378]), 295: array([614, 633, 636, 655, 662]), 296: array([ 755, 1240, 738, 1263, 364]), 297: array([1050, 1051, 1029, 1052, 1198]), 298: array([ 514, 1768, 527, 525, 1368]), 299: array([ 141, 935, 278, 80, 1193]), 300: array([1138, 1117, 1125, 370, 1114]), 301: array([1232, 1222, 1220, 1351, 1223]), 302: array([1660, 1645, 386, 1470, 109]), 303: array([1646, 461, 982, 1472, 1643]), 304: array([ 416, 425, 394, 1594, 1600]), 305: array([1469, 1466, 1478, 1059, 1060]), 306: array([271, 12, 313, 919, 33]), 307: array([184, 165, 172, 182, 181]), 308: array([756, 728, 255, 713, 762]), 309: array([121, 244, 746, 628, 59]), 310: array([439, 444, 401, 409, 433]), 311: array([660, 630, 636, 662, 654]), 312: array([1766, 1774, 1763, 399, 1400]), 313: array([322, 294, 355, 303, 324]), 314: array([ 816, 829, 1651, 869, 1451]), 315: array([ 888, 829, 1011, 1668, 461]), 316: array([ 589, 567, 297, 1473, 313]), 317: array([1318, 1237, 677, 1319, 1034]), 318: array([ 700, 1289, 1035, 1281, 1274]), 319: array([ 250, 369, 1524, 672, 701]), 320: array([ 700, 1289, 1035, 1281, 1274]), 321: array([1161, 382, 1166, 57, 127]), 322: array([414, 432, 400, 428, 398]), 323: array([382, 645, 701, 672, 93]), 324: array([838, 856, 842, 829, 841]), 325: array([1651, 1521, 261, 761, 816]), 326: array([ 979, 804, 996, 18, 1415]), 327: array([1144, 1354, 1295, 1017, 656]), 328: array([494, 481, 492, 505, 506]), 329: array([ 950, 656, 1185, 1180, 1413]), 330: array([398, 393, 404, 394, 432]), 331: array([1569, 521, 1559, 1540, 534]), 332: array([ 170, 184, 1662, 1010, 461]), 333: array([1049, 1028, 1044, 1051, 1047]), 334: array([1201, 1200, 1170, 1029, 1694]), 335: array([427, 436, 413, 442, 447]), 336: array([ 300, 1294, 844, 1651, 1044]), 337: array([1666, 12, 461, 982, 970]), 338: array([394, 422, 432, 412, 442]), 339: array([1235, 1252, 1034, 1240, 1218]), 340: array([1053, 1048, 1680, 1047, 957]), 341: array([506, 14, 137, 470, 492]), 342: array([1738, 791, 1766, 1584, 1363]), 343: array([1027, 1014, 1015, 1040, 1668]), 344: array([ 98, 1418, 339, 1295, 649]), 345: array([ 138, 529, 523, 1762, 1472]), 346: array([1108, 1102, 1103, 1067, 1062]), 347: array([1719, 1697, 1716, 1720, 1718]), 348: array([ 554, 1363, 402, 1378, 433]), 349: array([175, 110, 101, 270, 4]), 350: array([1627, 1619, 1596, 1594, 1621]), 351: array([ 619, 655, 1035, 652, 1133]), 352: array([1236, 1211, 1251, 58, 1209]), 353: array([1009, 229, 970, 196, 274]), 354: array([ 743, 761, 274, 996, 1010]), 355: array([1261, 1259, 1268, 1277, 1264])}\n",
"Precision@1 on Validation Set: 0.0\n"
]
}
],
"source": [
"# User\n",
"# valid_documents = valid['Description'].tolist()\n",
"# results = {}\n",
"\n",
"# for i, query in enumerate(valid_documents):\n",
"# query_embedding = model.encode([query], convert_to_tensor=True)\n",
"# distances, indices = index.search(query_embedding.cpu().numpy(), k=5)\n",
"# top_indices = indices[0]\n",
"# results[i] = top_indices\n",
"\n",
"# def evaluate(results):\n",
"# relevant_count = 0\n",
"# total_queries = len(results)\n",
"# for i, (_, top_indices) in enumerate(results.items()):\n",
"# # Check if any of the top indices match the index of the query document in the validation set\n",
"# if i in top_indices[:1]:\n",
"# relevant_count += 1\n",
"# precision_at_1 = relevant_count / total_queries\n",
"# return precision_at_1\n",
"\n",
"\n",
"# precision_at_1 = evaluate(results)\n",
"# print(\"Precision@1 on Validation Set:\", precision_at_1)"
]
},
{
"cell_type": "code",
"execution_count": 49,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"START\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"Some weights of BertForSequenceClassification were not initialized from the model checkpoint at BAAI/bge-small-en-v1.5 and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
"You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"[ 152 1363 1372 1373 891]\n",
"Rank 1: Starship Troopers (Score: 35.16937255859375)\n",
"Rank 2: Gods from Outer Space (Score: 35.109352111816406)\n",
"Rank 3: Voodoo Science: The Road from Foolishness to Fraud (Score: 34.12990951538086)\n"
]
}
],
"source": [
"def semantic_search(query, k=5, rerank_k=3, flag_reranker=None, flag_threshold=0):\n",
" \n",
" query_embedding = sbert_model.encode([query], convert_to_tensor=True)\n",
" distances, indices = index.search(query_embedding.cpu().numpy(), k + rerank_k)\n",
"\n",
" initial_indices = indices[0][:k]\n",
"\n",
" initial_documents = df.iloc[initial_indices][['Title', 'Description', 'Genres']]\n",
" genres_text = ''.join(initial_documents['Genres'].to_list())\n",
"\n",
" initial_documents['Text'] = initial_documents['Title'].str.lower() + ' ' + initial_documents['Description'].str.lower() + genres_text # Concatenate title and description\n",
" initial_distances = distances[0][:k]\n",
"\n",
" initial_results = list(zip(initial_documents['Title'], initial_documents['Text'], initial_distances))\n",
"\n",
" if flag_reranker:\n",
" flag_scores = [flag_reranker.compute_score([query, text]) for _, text, _ in initial_results]\n",
" reranked_results = [(title, text, dist + flag_score) for title, text, dist, flag_score in zip(initial_documents['Title'], initial_documents['Text'], initial_distances, flag_scores) if abs(flag_score) > flag_threshold]\n",
" reranked_results = sorted(reranked_results, key=lambda x: x[2], reverse=True)[:rerank_k]\n",
" else:\n",
" reranked_results = initial_results[:rerank_k]\n",
"\n",
" return reranked_results\n",
"\n",
"print(\"START\")\n",
"query = \"Search for science fiction books with a focus on space exploration.\"\n",
"flag_reranker = FlagReranker('BAAI/bge-small-en-v1.5', use_fp16=True) # Initialize the FlagReranker object\n",
"results = semantic_search(query, flag_reranker=flag_reranker)\n",
"\n",
"for rank, (title, text, score) in enumerate(results, start=1):\n",
" print(f\"Rank {rank}: {title} (Score: {score})\")\n"
]
}
],
"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.9.6"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
|