updates

app.py

@@ -7,13 +7,13 @@ import pyterrier as pt
 pt.init()
 import pyt_splade
 from pyterrier_gradio import Demo, MarkdownFile, interface, df2code, code2md, EX_Q, EX_D
-factory_max = pyt_splade.SpladeFactory(agg='max')
-factory_sum = pyt_splade.SpladeFactory(agg='sum')
+factory_max = pyt_splade.Splade(agg='max')
+factory_sum = pyt_splade.Splade(agg='sum')
 
 COLAB_NAME = 'pyterrier_splade.ipynb'
 COLAB_INSTALL = '''
 !pip install -q git+https://github.com/naver/splade
-!pip install -q git+https://github.com/seanmacavaney/pyt_splade@misc
+!pip install -q git+https://github.com/cmacdonald/pyt_splade
 '''.strip()
 
 def generate_vis(df, mode='Document'):
@@ -24,15 +24,9 @@ def generate_vis(df, mode='Document'):
     max_score = max(max(t.values()) for t in df['toks'])
   for row in df.itertuples(index=False):
     if mode == 'Query':
-      tok_scores = {m.group(2): float(m.group(1)) for m in re.finditer(r'#combine:0=([0-9.]+)\((#base64\([^)]+\)|[^)]+)\)', row.query)}
-      for key, value in list(tok_scores.items()):
-        if key.startswith('#base64('):
-          b64 = re.search('#base64\(([^)]+)\)', key).group(1)
-          del tok_scores[key]
-          key = base64.b64decode(b64).decode()
-          tok_scores[key] = value
+      tok_scores = row.query_toks
+      orig_tokens = factory_max.tokenizer.tokenize(row.text)
       max_score = max(tok_scores.values())
-      orig_tokens = factory_max.tokenizer.tokenize(row.query_0)
       id = row.qid
     else:
       tok_scores = row.toks
@@ -55,38 +49,36 @@ def generate_vis(df, mode='Document'):
 
 def predict_query(input, agg):
   code = f'''import pandas as pd
-import pyterrier as pt ; pt.init()
 import pyt_splade
 
-splade = pyt_splade.SpladeFactory(agg={repr(agg)})
+splade = pyt_splade.Splade(agg={agg!r})
 
-query_pipeline = splade.query()
+query_pipeline = splade.query_encoder()
 
 query_pipeline({df2code(input)})
 '''
   pipeline = {
     'max': factory_max,
     'sum': factory_sum
-  }[agg].query()
+  }[agg].query_encoder()
   res = pipeline(input)
   vis = generate_vis(res, mode='Query')
   return (res, code2md(code, COLAB_INSTALL, COLAB_NAME), vis)
 
 def predict_doc(input, agg):
   code = f'''import pandas as pd
-import pyterrier as pt ; pt.init()
 import pyt_splade
 
-splade = pyt_splade.SpladeFactory(agg={repr(agg)})
+splade = pyt_splade.Splade(agg={repr(agg)})
 
-doc_pipeline = splade.indexing()
+doc_pipeline = splade.doc_encoder()
 
 doc_pipeline({df2code(input)})
 '''
   pipeline = {
     'max': factory_max,
     'sum': factory_sum
-  }[agg].indexing()
+  }[agg].doc_encoder()
   res = pipeline(input)
   vis = generate_vis(res, mode='Document')
   res['toks'] = [json.dumps({k: round(v, 4) for k, v in t.items()}) for t in res['toks']]