diff --git a/wani.py b/wani.py
index 23710f2..d1e3db3 100644
--- a/wani.py
+++ b/wani.py
@@ -782,6 +782,31 @@ def load_corpus(filename):
 
 def main():
     import time
+class ColocationIds:
+    def __init__(self):
+        self.data = {}
+
+    def add_match(self, key):
+        if key in self.data:
+            self.data[key][1] += 1
+        else:
+            self.data[key] = [str(len(self.data) + 1), 1, False]
+    
+    def get(self, key, n):
+        return self.data[key][n]
+
+    def is_written(self, key):
+        return self.get(key, 2)
+
+    def num(self, key):
+        return str(self.get(key, 1))
+
+    def to_id(self, key):
+        return self.get(key, 0)
+
+    def set_written(self, key):
+        self.data[key][2] = True
+
     t = time.time()
 
     structures = build_structures(structures_file)