Research-Stack/2-Search-Space/search/stract/ltr/leechy_annotate.py

91 lines
2.1 KiB
Python

import leechy
from db import Db
import stract
import numpy as np
from tqdm import tqdm
import time
import sys
NUM_LABELS = 4
with open("data/queries_us.csv") as f:
all_queries = []
for query in [line.strip() for line in f.readlines()]:
if len(query) < 3:
continue
# check if query has large percentage of non-alphanumeric characters
if sum([c.isalnum() for c in query]) / len(query) < 0.5:
continue
if len(query) > 100:
continue
if len(query.split()) <= 1:
continue
all_queries.append(query)
# shuffle queries
np.random.shuffle(all_queries)
db = Db("data/auto-ranking-annotation.sqlite")
for query in all_queries:
db.add_query(query)
unannotated_queries = db.get_unannotated_queries()
eng = leechy.Engine()
for qid, query in tqdm(unannotated_queries.items()):
if query not in all_queries:
continue
tqdm.write(query)
leechy_results = []
for i in range(3):
leechy_results = eng.search(query)
if len(leechy_results) > 0:
break
if len(leechy_results) == 0:
tqdm.write(f"No results found for {query}")
sys.exit(-1)
for i, result in enumerate(leechy_results):
label = NUM_LABELS - int(np.log2(i + 1))
try:
data = next(
res
for res in stract.search(f"{query} exacturl:{result}")
if res["url"] == result
)
except StopIteration:
continue
if not data:
continue
tqdm.write(f"{result}: {label}")
db.insert_result(qid, i, data)
db.annotate(qid, result, label)
for page in range(2, 2 + 1):
bad_results = stract.search(query, page=page)
for i, result in enumerate(bad_results):
if result["url"] in leechy_results:
continue
label = 0
tqdm.write(f'{result["url"]}: {label}')
db.insert_result(qid, (i + 1) * page, result)
db.annotate(qid, result["url"], label)
time.sleep(np.random.normal(60, 20))
tqdm.write("")