User-friendly API and README

README.md

@@ -0,0 +1,48 @@
+## What is it?
+It's a UD segmenter.  
+More on UD [here](https://universaldependencies.org/).
+
+
+## How to use it?
+### PyPy
+It is highly recommended to use [PyPy](https://www.pypy.org/) istead of standard CPython interpreter.  
+To install al necessary third-party libraries for PyPy, run the following commands:
+```
+wget https://bootstrap.pypa.io/get-pip.py
+pypy3 get-pip.py --user
+pypy3 -m pip install distance --user
+pypy3 -m pip install tqdm --user
+```
+### Data
+The segmenter accepts data in lemma-tag format. Example input:
+```
+avata number=singv person=3 avaa
+hauska case=par number=plur hauskoja
+erikoinen case=par number=plur erikoisia
+...
+```
+Example output:
+```
+avaa/avata >/person=3,number=singv
+hausko/hauska j/number=plur a>/case=par
+erikois/erikoinen i/number=plur a>/case=par
+...
+```
+
+### Script
+Example usage:
+```
+pypy3 ud_segmenter.py --train data/fitrain+test --test data/fitest -n 1000 -o fi_segmented
+```
+#### Arguments
+* `--train` — train dataset file in lemma-tags form.
+* `--test` — test dataset file in lemma-tags form.
+* `-n` or `--number_of_passes` — number of passes for the model.
+* `-o` or `--output` — output file for the result of the segmentation.
+
+### Module
+Instead of running the program as a script you could use it as a module.  
+```python
+import ud_segmenter
+segmented = ud_segmenter.run(train_data, test_data, N=1000)
+```

base_sample.py

@@ -3,6 +3,7 @@
 from random import random, shuffle, seed
 from math import log
 from distance import levenshtein
+from tqdm import tqdm
 
 from split import sample_next, sample_nth, get_init_split, get_str
 
@@ -141,132 +142,130 @@ def filter_labels(labels):
 #                                          "num=PLV","num=SGV","case=NOM"]]
 
 
-train_fn = argv[1]
-test_fn = argv[2]
-
-train_data =  open(train_fn).read().split('\n')
-shuffle(train_data)
-test_data =  open(test_fn).read().split('\n')
-
-all_data = train_data + test_data
-all_data = [l.split(' ') for l in all_data if l != '']
-
-base_jcounts = defaultdict(lambda : defaultdict(lambda : 0.0))
-base_counts = defaultdict(lambda : 0.0)
-jcounts = defaultdict(lambda : defaultdict(lambda : 0.0))
-counts = defaultdict(lambda : 0.0)
-tcounts = defaultdict(lambda : defaultdict(lambda : 1.0))
-
-params = (base_jcounts,
-          base_counts,
-          jcounts,
-          counts,
-          tcounts)
-all_labels = set()
-for fields in all_data:
-    wf = fields[-1] + ">"
-    fields = filter_labels(fields[:len(fields) - 1])
-    all_labels.update(fields)
-    for i in range(len(wf)):
-        for j in range(i+1, len(wf) + 1):
-            base_counts[wf[i:j]] += 1
-            for l in fields:
-                base_jcounts[wf[i:j]][gl(l)] += 1
-#                base_counts[wf[i:j]] += 1
-#                base_counts[gl(l)] += 1
-#    for l in fields:
-#        base_jcounts[""][gl(l)] += 1
-#        base_counts[""] += 1
-#        base_counts[gl(l)] += 1
-
-for l in all_labels:
-    base_jcounts[""][gl(l)] += 1
-    base_counts[""] += 1
-    base_counts[gl(l)] += 1
-
-splits = []
-assignments = []
-labels = []
-wfs = []
-
-for fields in all_data:
-    wf =  fields[-1] + ">"
-    fields = filter_labels(fields[:len(fields) - 1])
-    wfs.append(wf)
-    labels.append(fields)
-    s = get_init_split(wf)
-    s = sample_nth(100,s,len(fields))
-    assignment = get_assignment(s,fields,wf,ALPHA,params)
-    splits.append(s)
-    assignments.append(assignment)
-    update_counts(s,assignment,wf,params,1)
-
-N = 1000
-
-best_splits = []
-best_assignment = []
-best_H = float('inf')
-logfile = open("logfile",'w')
-for n in range(N):
-    for j, s in enumerate(splits):
-        new_split = sample_next(splits[j],len(labels[j]))
-        new_assignment = get_assignment(new_split, labels[j], wfs[j], ALPHA, params)
-        old_prob = get_prob(splits[j], assignments[j],wfs[j],ALPHA,params)
-        new_prob = get_prob(new_split, new_assignment,wfs[j],ALPHA,params)
-        if (new_prob / old_prob)**2 > random():
-#        if new_prob > old_prob:
-            update_counts(splits[j],assignments[j],wfs[j],params,-1)
-            splits[j] = new_split
-            assignments[j] = new_assignment
-            update_counts(splits[j],assignments[j],wfs[j],params,1)
-    H = entropy(params)
-    logfile.write("%.3f\n" % H)
-    logfile.flush()
-    if H < best_H:
-        best_splits = list(splits)
-        best_assignments = list(assignments)
-        best_H = H
-    stderr.write("%u of passes %u, %.3f\r"# '>':tense=fut %f\r"#, ssa:ine %f isi:cond %f i:pln %f \r" % 
-                 % (n+1,N,entropy(params),
-#                 params[JC][">"][gl("tense=fut")]/params[C][gl("tense=fut")],
-#                  params[JC]["isi"][gl("mood=cnd")]/params[C][gl("mood=cnd")],
-#                  params[JC]["i"][gl("number=plur")]/params[C][gl("number=plur")],)
-                    ))
-
-
-stderr.write('\n')
-ssac = 0
-inc = 0
-
-for i,a  in enumerate(best_assignments):
-    if i + 1 < len(train_data):
-        continue
-    best_split = best_splits[i]
-    best_assign = best_assignments[i]
-    best_score = get_prob(best_split, best_assign,wfs[i],ALPHA,params)
-
-    split = best_split
-    assign = best_assign
-
-    for k in range(5000):
-        old_score = get_prob(split, assign,wfs[i],ALPHA,params)
-        new_split = sample_next(split,len(labels[i]) - 1)
-        new_assign = get_assignment(new_split, labels[i], wfs[i], ALPHA, params)
-        new_score = get_prob(new_split, new_assign,wfs[i],ALPHA,params)
-
-        if (new_score/old_score)**2 > random():
-            split = new_split
-            assign = new_assign
-        if new_score > best_score:
-            best_split = split
-            best_assign = assign
-            best_score = new_score
-
-    ss = [get_str(j,best_split,wfs[i]) for j in range(len(best_split) - 1)]
-#    ss = [get_str(j,best_split,wfs[i]) for j in range(len(best_split))]
-    for j, s in enumerate(ss):
-        stdout.write("%s/%s"%(s,','.join(best_assign[j])))
-        if j + 1 < len(ss):
-            stdout.write(" ")
-    print()
-    stdout.flush()
+def segment(train_data, test_data, N=1000):
+    shuffle(train_data)
+
+    all_data = train_data + test_data
+    all_data = [l.split(' ') for l in all_data if l != '']
+
+    base_jcounts = defaultdict(lambda : defaultdict(lambda : 0.0))
+    base_counts = defaultdict(lambda : 0.0)
+    jcounts = defaultdict(lambda : defaultdict(lambda : 0.0))
+    counts = defaultdict(lambda : 0.0)
+    tcounts = defaultdict(lambda : defaultdict(lambda : 1.0))
+
+    params = (base_jcounts,
+            base_counts,
+            jcounts,
+            counts,
+            tcounts)
+    all_labels = set()
+    for fields in all_data:
+        wf = fields[-1] + ">"
+        fields = filter_labels(fields[:len(fields) - 1])
+        all_labels.update(fields)
+        for i in range(len(wf)):
+            for j in range(i+1, len(wf) + 1):
+                base_counts[wf[i:j]] += 1
+                for l in fields:
+                    base_jcounts[wf[i:j]][gl(l)] += 1
+    #                base_counts[wf[i:j]] += 1
+    #                base_counts[gl(l)] += 1
+    #    for l in fields:
+    #        base_jcounts[""][gl(l)] += 1
+    #        base_counts[""] += 1
+    #        base_counts[gl(l)] += 1
+
+    for l in all_labels:
+        base_jcounts[""][gl(l)] += 1
+        base_counts[""] += 1
+        base_counts[gl(l)] += 1
+
+    splits = []
+    assignments = []
+    labels = []
+    wfs = []
+
+    for fields in all_data:
+        wf =  fields[-1] + ">"
+        fields = filter_labels(fields[:len(fields) - 1])
+        wfs.append(wf)
+        labels.append(fields)
+        s = get_init_split(wf)
+        s = sample_nth(100,s,len(fields))
+        assignment = get_assignment(s,fields,wf,ALPHA,params)
+        splits.append(s)
+        assignments.append(assignment)
+        update_counts(s,assignment,wf,params,1)
+
+    best_splits = []
+    best_assignment = []
+    best_H = float('inf')
+    logfile = open("logfile",'w')
+    for n in range(N):
+        for j, s in enumerate(splits):
+            new_split = sample_next(splits[j],len(labels[j]))
+            new_assignment = get_assignment(new_split, labels[j], wfs[j], ALPHA, params)
+            old_prob = get_prob(splits[j], assignments[j],wfs[j],ALPHA,params)
+            new_prob = get_prob(new_split, new_assignment,wfs[j],ALPHA,params)
+            if (new_prob / old_prob)**2 > random():
+    #        if new_prob > old_prob:
+                update_counts(splits[j],assignments[j],wfs[j],params,-1)
+                splits[j] = new_split
+                assignments[j] = new_assignment
+                update_counts(splits[j],assignments[j],wfs[j],params,1)
+        H = entropy(params)
+        logfile.write("%.3f\n" % H)
+        logfile.flush()
+        if H < best_H:
+            best_splits = list(splits)
+            best_assignments = list(assignments)
+            best_H = H
+        stderr.write("%u of passes %u, %.3f\r"# '>':tense=fut %f\r"#, ssa:ine %f isi:cond %f i:pln %f \r" % 
+                    % (n+1,N,entropy(params),
+    #                 params[JC][">"][gl("tense=fut")]/params[C][gl("tense=fut")],
+    #                  params[JC]["isi"][gl("mood=cnd")]/params[C][gl("mood=cnd")],
+    #                  params[JC]["i"][gl("number=plur")]/params[C][gl("number=plur")],)
+                        ))
+
+
+    stderr.write('\n')
+    ssac = 0
+    inc = 0
+
+    result = []
+    with tqdm(total=len(best_assignments)-len(train_data)) as pbar:
+        for i,a  in enumerate(best_assignments):
+            if i + 1 < len(train_data):
+                continue
+            best_split = best_splits[i]
+            best_assign = best_assignments[i]
+            best_score = get_prob(best_split, best_assign,wfs[i],ALPHA,params)
+
+            split = best_split
+            assign = best_assign
+
+            for k in range(5000):
+                old_score = get_prob(split, assign,wfs[i],ALPHA,params)
+                new_split = sample_next(split,len(labels[i]) - 1)
+                new_assign = get_assignment(new_split, labels[i], wfs[i], ALPHA, params)
+                new_score = get_prob(new_split, new_assign,wfs[i],ALPHA,params)
+
+                if (new_score/old_score)**2 > random():
+                    split = new_split
+                    assign = new_assign
+                if new_score > best_score:
+                    best_split = split
+                    best_assign = assign
+                    best_score = new_score
+
+            ss = [get_str(j,best_split,wfs[i]) for j in range(len(best_split) - 1)]
+            temp = []
+        #    ss = [get_str(j,best_split,wfs[i]) for j in range(len(best_split))]
+            for j, s in enumerate(ss):
+                temp.append("%s/%s"%(s,','.join(best_assign[j])))
+                if j + 1 < len(ss):
+                    temp.append(" ")
+            result.append(''.join(temp))
+            pbar.update(1)
+    return '\n'.join(result)

ud_segmenter.py

@@ -0,0 +1,63 @@
+#!/usr/bin/pypy3
+import base_sample
+
+def run(train_data, test_data, N=1000):
+    """Runs the model on given datasets and outputs the segmented version
+
+    Parameters
+    ----------
+    train_data: str
+        A string containing train data in lemma-tags form.
+    test_data: str
+        Same for test data.
+    N: int
+        Number of passes.
+    """
+    train_data = train_data.split('\n')
+    test_data = test_data.split('\n')
+    return base_sample.segment(train_data, test_data, N=N)
+
+if __name__ == '__main__':
+    import argparse
+
+    arg_parser = argparse.ArgumentParser()
+    arg_parser.add_argument(
+        '--train',
+        metavar='filename',
+        type=str,
+        required=True,
+        help='Train dataset file in lemma-tags form.',
+    )
+    arg_parser.add_argument(
+        '--test',
+        metavar='filename',
+        type=str,
+        required=True,
+        help='Test dataset file in lemma-tags form.',
+    )
+    arg_parser.add_argument(
+        '-o',
+        '--output',
+        metavar='filename',
+        type=str,
+        required=True,
+        help='Output file for the result of the segmentation.',
+    )
+    arg_parser.add_argument(
+        '-n',
+        '--number_of_passes',
+        metavar='N',
+        type=int,
+        required=True,
+        default=1000,
+        help='Number of passes for the model.',
+    )
+    args = arg_parser.parse_args()
+
+    with open(args.train) as f:
+        train_data = f.read()
+    with open(args.test) as f:
+        test_data = f.read()
+    result = run(train_data, test_data, N=args.number_of_passes)
+    with open(args.output, 'w') as f:
+        f.write(result)