70 lines
2.0 KiB
Python
70 lines
2.0 KiB
Python
import csv
|
|
import os
|
|
import requests
|
|
from time import sleep
|
|
|
|
HEADERS = {"User-Agent": "cse881"}
|
|
SEARCH_URL = "https://en.wikipedia.org/w/api.php"
|
|
BASE_URL = "https://en.wikipedia.org/api/rest_v1"
|
|
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
INPUT_TSV = os.path.abspath(os.path.join(BASE_DIR, "../data/raw/imdb_datasets/title.basics.test.tsv"))
|
|
OUTPUT_DIR = os.path.abspath(os.path.join(BASE_DIR, "../data/raw/wikipedia/wikipedia_html"))
|
|
|
|
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
|
|
|
def fetch_wikipedia_html(query):
|
|
params = {
|
|
"action": "query",
|
|
"list": "search",
|
|
"srsearch": query,
|
|
"format": "json"
|
|
}
|
|
|
|
resp = requests.get(SEARCH_URL, params=params, headers=HEADERS).json()
|
|
results = resp.get("query", {}).get("search", [])
|
|
|
|
if not results:
|
|
return None
|
|
|
|
best_title = results[0]["title"]
|
|
wiki_title = best_title.replace(" ", "_")
|
|
html_url = f"{BASE_URL}/page/html/{wiki_title}"
|
|
r = requests.get(html_url, headers=HEADERS)
|
|
|
|
if r.status_code != 200:
|
|
return None
|
|
return r.text
|
|
|
|
|
|
with open(INPUT_TSV, encoding="utf-8") as f:
|
|
print("Opened file:", INPUT_TSV)
|
|
print("First 500 chars:")
|
|
print(f.read(500))
|
|
f.seek(0)
|
|
|
|
reader = csv.DictReader(f, delimiter="\t")
|
|
for row in reader:
|
|
tconst = row["tconst"]
|
|
title = row["primaryTitle"]
|
|
year = row["startYear"]
|
|
outfile = os.path.join(OUTPUT_DIR, f"{tconst}.html")
|
|
print(outfile)
|
|
|
|
if os.path.exists(outfile):
|
|
print(f"Skipping {tconst}: {query}")
|
|
continue #if exists, skip
|
|
|
|
query = f"{title} {year}" if year != "\\N" else title
|
|
print(f"Fetching Wikipedia for {tconst}: {query}")
|
|
html = fetch_wikipedia_html(query)
|
|
if html:
|
|
with open(outfile, "w", encoding="utf-8") as out:
|
|
out.write(html)
|
|
else:
|
|
print(f"No Wikipedia page found")
|
|
sleep(0.5)
|
|
print("Completed")
|
|
|
|
#https://en.wikipedia.org/w/index.php?api=wmf-restbase&title=Special%3ARestSandbox#/Page%20content/get_page_summary__title_
|