From f784d7bcd24d9befe4f82c23b77139b0b30a5f5b Mon Sep 17 00:00:00 2001
From: Piotr Czapla <github@piotrczapla.com>
Date: Tue, 1 Jan 2019 15:13:14 +0100
Subject: [PATCH] Make the article detection code work with our wikitext

---
 ulmfit/pretrain_lm.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/ulmfit/pretrain_lm.py b/ulmfit/pretrain_lm.py
index 27b3fb3..9ee559c 100644
--- a/ulmfit/pretrain_lm.py
+++ b/ulmfit/pretrain_lm.py
@@ -35,7 +35,7 @@ class Tokenizers(Enum):
     FASTAI='f'
 
 def istitle(line):
-    return len(re.findall(r'^ = [^=]* = $', line)) != 0
+    return len(re.findall(r'^ ?= [^=]* = ?$', line)) != 0
 
 def read_wiki_articles(filename):
     articles = []
@@ -48,6 +48,7 @@ def read_wiki_articles(filename):
             articles.append(current_article)
             current_article = ''
     articles.append(current_article)
+    print(f"Wiki text was split to {len(articles)} articles")
     return pd.DataFrame({'texts':np.array(articles)})
 
 @dataclass