Add debug mode

app.py

@@ -9,8 +9,13 @@ import shutil
 
 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 
+# To suppress the following warning:
+# huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
 HF_TOKEN_DOWNLOAD = os.environ['HF_TOKEN_DOWNLOAD']
 HF_TOKEN_UPLOAD = os.environ['HF_TOKEN_UPLOAD']
+MODE = os.environ['MODE']
 
 MODEL_NAME = 'liujch1998/cd-pi'
 DATASET_REPO_URL = "https://huggingface.co/datasets/liujch1998/cd-pi-dataset"
@@ -32,6 +37,8 @@ repo.git_pull()
 class Interactive:
     def __init__(self):
         self.tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=HF_TOKEN_DOWNLOAD)
+        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.linear.weight = torch.nn.Parameter(self.model.shared.weight[32099, :].unsqueeze(0)) # (1, D)
@@ -40,6 +47,13 @@ class Interactive:
         self.t = self.model.shared.weight[32097, 0].item()
 
     def run(self, statement):
+        if MODE == 'debug':
+            return {
+                'logit': 0.0,
+                'logit_calibrated': 0.0,
+                'score': 0.5,
+                'score_calibrated': 0.5,
+            }
         input_ids = self.tokenizer.batch_encode_plus([statement], return_tensors='pt', padding='longest').input_ids.to(device)
         with torch.no_grad():
             output = self.model(input_ids)
@@ -55,16 +69,10 @@ class Interactive:
             'score': score.item(),
             'score_calibrated': score_calibrated.item(),
         }
-        # return {
-        #     'logit': 0.0,
-        #     'logit_calibrated': 0.0,
-        #     'score': 0.5,
-        #     'score_calibrated': 0.5,
-        # }
 
 interactive = Interactive()
 
-def predict(statement):
+def predict(statement, model):
     result = interactive.run(statement)
     with open(DATA_PATH, 'a') as f:
         row = {
@@ -155,7 +163,7 @@ description = '''This is a demo for a commonsense statement verification model.
 
 gr.Interface(
     fn=predict,
-    inputs=[input_statement],
+    inputs=[input_statement, input_model],
     outputs=output,
     title="cd-pi Demo",
     description=description,