generate n-grams in order, return positions to caller

Author: larsmans

semanticizest/_util.py

@@ -2,14 +2,18 @@
 from six.moves.urllib.parse import quote
 
 
-def ngrams(lst, N):
+def ngrams_with_pos(lst, N):
     """Generate n-grams for 1 <= n <= N from lst."""
 
     join = " ".join
 
-    for n in xrange(N):
-        for start in xrange(len(lst) - n):
-            yield join(lst[start:start + n + 1])
+    for start in xrange(len(lst)):
+        for n in xrange(1, 1 + min(N, len(lst) - start)):
+            yield start, start + n, join(lst[start:start + n])
+
+def tosequence(x):
+    """Cast x to sequence. Returns x if at all possible."""
+    return x if isinstance(x, Sequence) else list(x)
 
 
 def url_from_title(title, wiki):

semanticizest/tests/test_util.py

@@ -1,8 +1,8 @@
 from collections import Counter
 
-from semanticizest._util import ngrams, url_from_title
+from semanticizest._util import ngrams, ngrams_with_pos, url_from_title
 
-from nose.tools import assert_equal, assert_true
+from nose.tools import assert_equal, assert_in, assert_true
 
 
 def test_ngrams():
@@ -15,6 +15,10 @@ def test_ngrams():
     assert_equal(set(ng), expected)
     assert_true(all(freq == 1 for freq in ng.values()))
 
+    with_pos = list(ngrams_with_pos(text, N=2))
+    assert_in((0, 2, 'Hello ,'), with_pos)
+    assert_in((1, 3, ', world'), with_pos)
+
 
 def test_url_from_title():
     """Test article title -> Wikipedia URL conversion."""