Update app.py

app.py

@@ -48,14 +48,13 @@ def generate_vis(df, mode='Document'):
   return '\n'.join(result)
 
 def predict_query(input, agg):
-  code = f'''import pandas as pd
-import pyt_splade
+  code = f'''import pyt_splade
 
 splade = pyt_splade.Splade(agg={agg!r})
 
 query_pipeline = splade.query_encoder()
 
-query_pipeline({df2code(input)})
+query_pipeline({df2list(input)})
 '''
   pipeline = {
     'max': factory_max,
@@ -63,17 +62,17 @@ query_pipeline({df2code(input)})
   }[agg].query_encoder()
   res = pipeline(input)
   vis = generate_vis(res, mode='Query')
+  res['query_toks'] = [json.dumps({k: round(v, 4) for k, v in t.items()}) for t in res['query_toks']]
   return (res, code2md(code, COLAB_INSTALL, COLAB_NAME), vis)
 
 def predict_doc(input, agg):
-  code = f'''import pandas as pd
-import pyt_splade
+  code = f'''import pyt_splade
 
 splade = pyt_splade.Splade(agg={repr(agg)})
 
 doc_pipeline = splade.doc_encoder()
 
-doc_pipeline({df2code(input)})
+doc_pipeline({df2list(input)})
 '''
   pipeline = {
     'max': factory_max,