WIP

app.py

@@ -40,7 +40,8 @@ class Interactive:
         if MODE == 'debug':
             return
         self.model = transformers.T5EncoderModel.from_pretrained(MODEL_NAME, use_auth_token=HF_TOKEN_DOWNLOAD, low_cpu_mem_usage=True, device_map='auto', torch_dtype='auto')
-        self.linear = torch.nn.Linear(self.model.shared.embedding_dim, 1, dtype=self.model.dtype).to(device)
+        self.model.D = self.model.shared.embedding_dim
+        self.linear = torch.nn.Linear(self.model.D, 1, dtype=self.model.dtype).to(device)
         self.linear.weight = torch.nn.Parameter(self.model.shared.weight[32099, :].unsqueeze(0)) # (1, D)
         self.linear.bias = torch.nn.Parameter(self.model.shared.weight[32098, 0].unsqueeze(0)) # (1)
         self.model.eval()
@@ -64,52 +65,124 @@ class Interactive:
             score = logit.sigmoid()
             score_calibrated = logit_calibrated.sigmoid()
         return {
+            'timestamp': datetime.datetime.now().strftime('%Y%m%d-%H%M%S'),
+            'statement': statement,
             'logit': logit.item(),
             'logit_calibrated': logit_calibrated.item(),
             'score': score.item(),
             'score_calibrated': score_calibrated.item(),
         }
 
+    def runs(self, statements):
+        if MODE == 'debug':
+            return [{
+                'logit': 0.0,
+                'logit_calibrated': 0.0,
+                'score': 0.5,
+                'score_calibrated': 0.5,
+            } for _ in statements]
+        tok = self.tokenizer.batch_encode_plus(statements, return_tensors='pt', padding='longest')
+        input_ids = tok.input_ids.to(device)
+        attention_mask = tok.attention_mask.to(device)
+        with torch.no_grad():
+            output = self.model(input_ids=input_ids, attention_mask=attention_mask)
+            last_indices = attention_mask.sum(dim=1, keepdim=True) - 1 # (B, 1)
+            last_indices = last_indices.unsqueeze(-1).expand(-1, -1, self.model.D) # (B, 1, D)
+            last_hidden_state = output.last_hidden_state.to(device) # (B, L, D)
+            hidden = last_hidden_state.gather(dim=1, index=last_indices).squeeze(1) # (B, D)
+            logits = self.linear(hidden).squeeze(-1) # (B)
+            logits_calibrated = logits / self.t
+            scores = logits.sigmoid()
+            scores_calibrated = logits_calibrated.sigmoid()
+        return [{
+            'timestamp': datetime.datetime.now().strftime('%Y%m%d-%H%M%S'),
+            'statement': statement,
+            'logit': logit.item(),
+            'logit_calibrated': logit_calibrated.item(),
+            'score': score.item(),
+            'score_calibrated': score_calibrated.item(),
+        } for statement, logit, logit_calibrated, score, score_calibrated in zip(statements, logits, logits_calibrated, scores, scores_calibrated)]
+
 interactive = Interactive()
 
-def predict(statement, do_save=True):
-    result = interactive.run(statement)
-    output = {
-        'True': result['score_calibrated'],
-        'False': 1 - result['score_calibrated'],
-    }
-    output_raw = {
-        'timestamp': datetime.datetime.now().strftime('%Y%m%d-%H%M%S'),
-        'statement': statement,
-    }
-    output_raw.update(result)
-    if do_save:
-        with open(DATA_PATH, 'a') as f:
-            json.dump(output_raw, f, ensure_ascii=False)
-            f.write('\n')
+# def predict(statement, do_save=True):
+#     output_raw = interactive.run(statement)
+#     output = {
+#         'True': output_raw['score_calibrated'],
+#         'False': 1 - output_raw['score_calibrated'],
+#     }
+#     if do_save:
+#         with open(DATA_PATH, 'a') as f:
+#             json.dump(output_raw, f, ensure_ascii=False)
+#             f.write('\n')
+#         commit_url = repo.push_to_hub()
+#         print('Logged statement to dataset:')
+#         print('Commit URL:', commit_url)
+#         print(output_raw)
+#         print()
+#     return output, output_raw, gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(value='Please provide your feedback before trying out another statement.')
+
+# def record_feedback(output_raw, feedback, do_save=True):
+#     if do_save:
+#         output_raw.update({ 'feedback': feedback })
+#         with open(DATA_PATH, 'a') as f:
+#             json.dump(output_raw, f, ensure_ascii=False)
+#             f.write('\n')
+#         commit_url = repo.push_to_hub()
+#         print('Logged feedback to dataset:')
+#         print('Commit URL:', commit_url)
+#         print(output_raw)
+#         print()
+#     return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(value='Thanks for your feedback! Now you can enter another statement.')
+# def record_feedback_agree(output_raw, do_save=True):
+#     return record_feedback(output_raw, 'agree', do_save)
+# def record_feedback_disagree(output_raw, do_save=True):
+#     return record_feedback(output_raw, 'disagree', do_save)
+
+def predict(statements, do_saves):
+    output_raws = interactive.runs(statements)
+    outputs = [{
+        'True': output_raw['score_calibrated'],
+        'False': 1 - output_raw['score_calibrated'],
+    } for output_raw in output_raws]
+    for output_raw, do_save in zip(output_raws, do_saves):
+        if do_save:
+            with open(DATA_PATH, 'a') as f:
+                json.dump(output_raw, f, ensure_ascii=False)
+                f.write('\n')
+    if any(do_saves):
         commit_url = repo.push_to_hub()
         print('Logged statement to dataset:')
         print('Commit URL:', commit_url)
         print(output_raw)
         print()
-    return output, output_raw, gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(value='Please provide your feedback before trying out another statement.')
-
-def record_feedback(output_raw, feedback, do_save=True):
-    if do_save:
-        output_raw.update({ 'feedback': feedback })
-        with open(DATA_PATH, 'a') as f:
-            json.dump(output_raw, f, ensure_ascii=False)
-            f.write('\n')
+    return outputs, output_raws, \
+        [gr.update(visible=False) for _ in statements], \
+        [gr.update(visible=True) for _ in statements], \
+        [gr.update(visible=True) for _ in statements], \
+        [gr.update(value='Please provide your feedback before trying out another statement.') for _ in statements]
+
+def record_feedback(output_raws, feedback, do_saves):
+    for output_raw, do_save in zip(output_raws, do_saves):
+        if do_save:
+            output_raw.update({ 'feedback': feedback })
+            with open(DATA_PATH, 'a') as f:
+                json.dump(output_raw, f, ensure_ascii=False)
+                f.write('\n')
+    if any(do_saves):
         commit_url = repo.push_to_hub()
         print('Logged feedback to dataset:')
         print('Commit URL:', commit_url)
         print(output_raw)
         print()
-    return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(value='Thanks for your feedback! Now you can enter another statement.')
-def record_feedback_agree(output_raw, do_save=True):
-    return record_feedback(output_raw, 'agree', do_save)
-def record_feedback_disagree(output_raw, do_save=True):
-    return record_feedback(output_raw, 'disagree', do_save)
+    return [gr.update(visible=True) for _ in output_raws], \
+        [gr.update(visible=False) for _ in output_raws], \
+        [gr.update(visible=False) for _ in output_raws], \
+        [gr.update(value='Thanks for your feedback! Now you can enter another statement.') for _ in output_raws]
+def record_feedback_agree(output_raws, do_saves):
+    return record_feedback(output_raws, 'agree', do_saves)
+def record_feedback_disagree(output_raws, do_saves):
+    return record_feedback(output_raws, 'disagree', do_saves)
 
 examples = [
     # openbookqa
@@ -223,9 +296,12 @@ with gr.Blocks() as demo:
                 cache_examples=False,
                 run_on_click=False, # If we want this to be True, I suspect we need to enable the statement.submit()
             )
-    submit.click(predict, inputs=[statement, do_save], outputs=[output, output_raw, submit, feedback_agree, feedback_disagree, feedback_ack])
+    submit.click(predict, inputs=[statement, do_save], outputs=[output, output_raw, submit, feedback_agree, feedback_disagree, feedback_ack], batch=True, max_batch_size=16)
     # statement.submit(predict, inputs=[statement], outputs=[output, output_raw])
-    feedback_agree.click(record_feedback_agree, inputs=[output_raw, do_save], outputs=[submit, feedback_agree, feedback_disagree, feedback_ack])
-    feedback_disagree.click(record_feedback_disagree, inputs=[output_raw, do_save], outputs=[submit, feedback_agree, feedback_disagree, feedback_ack])
+    feedback_agree.click(record_feedback_agree, inputs=[output_raw, do_save], outputs=[submit, feedback_agree, feedback_disagree, feedback_ack], batch=True, max_batch_size=16)
+    feedback_disagree.click(record_feedback_disagree, inputs=[output_raw, do_save], outputs=[submit, feedback_agree, feedback_disagree, feedback_ack], batch=True, max_batch_size=16)
 
 demo.queue(concurrency_count=16).launch(debug=True)
+
+# Concurrency, Batching
+# Theme, CSS