wiki-mad-libs/test.py

98 lines
2.9 KiB
Python

import json
import random
import re
import requests
import time
from nltk import pos_tag, sent_tokenize, word_tokenize
from nltk.corpus import stopwords
from wikipediaapi import Wikipedia, WikipediaPage
# Info about the default pos_tag tags
# https://www.ling.upenn.edu/courses/Fall_2003/ling001/penn_treebank_pos.html
adlib_tags = {
"JJ": "Adjective",
"JJR": "Adjective ending in 'er'",
"JJS": "Adjective ending in 'est'",
"NN": "Noun",
"NNS": "Plural Noun",
"NNP": "Proper Noun",
"NNPS": "Plural Proper Noun",
"RB": "Adverb",
"RBR": "Adverb ending in 'er'",
"RBS": "Adverb ending in 'est'",
"VB": "Verb",
"VBD": "Past Tense Verb",
"VBG": "Verb ending in 'ing'",
"VBN": "Past Tense Verb",
"VBP": "Present Tense Verb",
"VBZ": "Present Tense Verb ending in 's'",
}
stop_words = set(stopwords.words("english"))
months = {
"january",
"february",
"march",
"april",
"may",
"june",
"july",
"august",
"september",
"october",
"november",
"december",
}
stop_words.update(months)
def get_random_wikipedia_title() -> str:
random_result = json.loads(requests.get('https://en.wikipedia.org/api/rest_v1/page/random/title').text)
return random_result['items'][0]['title']
def get_random_wikipedia_page(wikipedia: Wikipedia, min_length: int = None) -> WikipediaPage:
page = None
while(page is None):
page = wikipedia.page(get_random_wikipedia_title())
if(min_length and len(page.summary) < min_length):
print(f"{page.displaytitle} is too short. Retrying...")
page = None
time.sleep(1)
return page
wikipedia = Wikipedia('en')
wiki_page = get_random_wikipedia_page(wikipedia, 500)
print(wiki_page.title)
print(wiki_page.displaytitle)
print(wiki_page.canonicalurl)
summary = re.sub(r"(\w+)\.([A-Z])", r"\1. \2", wiki_page.summary) # Try to break up sentences mashed together by stripping strings
title = wiki_page.title
sentences = sent_tokenize(summary)
tagged_sentences = []
for sentence in sentences:
tagged_sentences.append(pos_tag(word_tokenize(sentence)))
i = adlib_word_counter = 0
min_words = 4
output_tokens = []
for sentence in tagged_sentences:
for token, tag in sentence:
output_tokens.append({"id": i, "token": token, "tag": tag})
adlib_tag = adlib_tags.get(tag)
if adlib_tag is not None:
if random.randint(0, adlib_word_counter) > min_words \
and token.lower() not in stop_words \
and token.lower() not in title.lower():
output_tokens[-1]["adlib_tag"] = adlib_tag
adlib_word_counter = 0
else:
adlib_word_counter += 1
print(f"{token}: {adlib_tag} {adlib_word_counter}")
i += 1
with open("article.js", "w") as json_file:
json_file.write(f"article = {json.dumps(output_tokens)}")