modify stat

docs/stat.py

@@ -3,6 +3,7 @@
 import glob
 import re
 
+import numpy as np
 import titlecase
 
 files = sorted(glob.glob('*_models.md'))
@@ -18,9 +19,10 @@
 
     # count papers
     papers = set(
-        titlecase.titlecase(x.lower().strip())
-        for x in re.findall(r'\btitle\s*=\s*{(.*)}', content))
-    paperlist = '\n'.join(sorted('    - ' + x for x in papers))
+        (papertype, titlecase.titlecase(paper.lower().strip()))
+        for (papertype, paper) in re.findall(
+            r'\[([A-Z]+?)\].*?\btitle\s*=\s*{(.*?)}', content, re.DOTALL))
+    paperlist = '\n'.join(sorted(f'    - [{t}] {x}' for t, x in papers))
 
     # count configs
     configs = set(x.lower().strip()
@@ -48,12 +50,18 @@
 allckpts = func.reduce(lambda a, b: a.union(b), [c for _, _, c, _ in stats])
 msglist = '\n'.join(x for _, _, _, x in stats)
 
+papertypes, papercounts = np.unique([t for t, _ in allpapers],
+                                    return_counts=True)
+countstr = '\n'.join(
+    [f'   - {t}: {c}' for t, c in zip(papertypes, papercounts)])
+
 modelzoo = f"""
 # Model Zoo Statistics
 
 * Number of checkpoints: {len(allckpts)}
 * Number of configs: {len(allconfigs)}
 * Number of papers: {len(allpapers)}
+{countstr}
 
 {msglist}
 """