stenodojo

sentence_generator.py (5772B)

---

 """Generate practice sentences from a constrained word pool.
 
 Uses NLTK's Brown and Gutenberg corpora — filters real English sentences
 to those where every word belongs to the target pool. Falls back to
 simple templates when the pool is too small to find corpus matches.
 """
 
 import re
 import random
 
 _corpus_cache = None
 
 
 def _normalize(word):
     """Lowercase, strip non-alpha except apostrophes (for contractions)."""
     return re.sub(r"[^a-z']", "", word.lower())
 
 
 def _load_corpus():
     """Lazily load and pre-process NLTK sentences into (word_set, text) pairs."""
     global _corpus_cache
     if _corpus_cache is not None:
         return _corpus_cache
 
     from nltk.corpus import brown, gutenberg
 
     raw = list(brown.sents()) + list(gutenberg.sents())
     corpus = []
     for tokens in raw:
         clean = _rebuild(tokens)
         if not clean:
             continue
         words = {_normalize(w) for w in clean.split()}
         words.discard("")
         if len(words) < 3:
             continue
         corpus.append((words, clean))
 
     _corpus_cache = corpus
     return corpus
 
 
 def _rebuild(tokens):
     """Reconstruct a sentence from NLTK tokens into clean readable text."""
     text = " ".join(tokens)
     # Fix NLTK tokenization artifacts
     text = re.sub(r" ([.,;:!?)\]}])", r"\1", text)
     text = re.sub(r"([\[({]) ", r"\1", text)
     text = text.replace("`` ", '"').replace(" ''", '"').replace("''", '"')
     text = text.replace('``', '"')
     text = text.replace(" n't", "n't").replace(" 're", "'re")
     text = text.replace(" 've", "'ve").replace(" 'll", "'ll")
     text = text.replace(" 'd", "'d").replace(" 's", "'s")
     text = text.replace(" 'm", "'m")
     # Strip leading/trailing quotes and whitespace
     text = text.strip().strip('"').strip()
     # Drop sentences with numbers
     if re.search(r"\d", text):
         return ""
     # Must start with a letter and end with sentence-ending punctuation
     if not text or not text[0].isalpha():
         return ""
     if not text.endswith((".", "!", "?")):
         return ""
     # Ensure first character is uppercase
     text = text[0].upper() + text[1:]
     # Skip very short or very long
     word_count = len(text.split())
     if word_count < 5 or word_count > 18:
         return ""
     # Skip sentences with remaining quote marks (dialogue fragments)
     if '"' in text or '``' in text or "''" in text:
         return ""
     return text
 
 
 def _filter_corpus(word_pool, count):
     """Return up to `count` corpus sentences using only words from pool."""
     corpus = _load_corpus()
     pool = {w.lower() for w in word_pool}
     # Also include common contractions if their base is in the pool
     extras = set()
     for w in list(pool):
         for suffix in ("n't", "'s", "'re", "'ve", "'ll", "'d", "'m"):
             extras.add(w + suffix)
     pool |= extras
 
     matches = []
     for word_set, text in corpus:
         if word_set.issubset(pool):
             matches.append(text)
 
     random.shuffle(matches)
     return matches[:count]
 
 
 # ── Fallback template generator for very small pools ──
 
 _PRONOUNS = {"i", "he", "she", "we", "they", "you", "it"}
 _VERBS = {
     "be", "have", "do", "say", "go", "get", "make", "know", "take", "come",
     "see", "think", "look", "want", "give", "use", "find", "tell", "ask",
     "work", "seem", "feel", "try", "leave", "call", "keep", "let", "begin",
     "run", "read", "show", "turn", "play", "learn", "set", "change", "move",
     "put", "start", "need", "help", "talk", "open", "mean", "add", "live",
 }
 _NOUNS = {
     "time", "people", "year", "day", "way", "man", "world", "life", "hand",
     "part", "place", "thing", "child", "eye", "woman", "work", "case",
     "point", "home", "water", "room", "mother", "area", "money", "story",
     "fact", "month", "lot", "right", "study", "book", "job", "word", "side",
     "head", "house", "name", "end", "door", "car", "food", "night", "state",
 }
 _DETS = {"the", "a", "an", "this", "that", "some", "any", "no"}
 _PREPS = {"in", "on", "at", "to", "for", "with", "from", "by", "about", "into", "over"}
 _AUXS = {"will", "would", "can", "could", "should", "must", "may", "might"}
 
 _TEMPLATES = [
     ("SUBJ", "VERB", "DET", "NOUN"),
     ("SUBJ", "AUX", "VERB"),
     ("SUBJ", "VERB", "PREP", "DET", "NOUN"),
     ("DET", "NOUN", "VERB", "PREP", "DET", "NOUN"),
     ("SUBJ", "AUX", "VERB", "DET", "NOUN"),
 ]
 
 _SLOT_MAP = {
     "SUBJ": _PRONOUNS, "VERB": _VERBS, "NOUN": _NOUNS,
     "DET": _DETS, "PREP": _PREPS, "AUX": _AUXS,
 }
 
 
 def _template_sentences(word_pool, count):
     pool = {w.lower() for w in word_pool}
     available = {}
     for slot, candidates in _SLOT_MAP.items():
         words = list(pool & candidates)
         if words:
             available[slot] = words
 
     usable = [t for t in _TEMPLATES if all(s in available for s in t)]
     if not usable:
         return []
 
     results = set()
     for _ in range(count * 10):
         pattern = random.choice(usable)
         words = [random.choice(available[s]) for s in pattern]
         words[0] = words[0].capitalize()
         words = ["I" if w == "i" else w for w in words]
         sent = " ".join(words) + "."
         results.add(sent)
         if len(results) >= count:
             break
     return list(results)[:count]
 
 
 def generate_sentences(word_pool, count=10):
     """Generate practice sentences constrained to word_pool.
 
     Tries corpus filtering first; falls back to templates for small pools.
     """
     if not word_pool:
         return []
 
     sentences = _filter_corpus(word_pool, count)
     if len(sentences) >= count:
         return sentences
 
     # Supplement with template-generated sentences
     needed = count - len(sentences)
     sentences.extend(_template_sentences(word_pool, needed))
     return sentences[:count]