bf16

app.py

@@ -32,35 +32,35 @@ repo.git_pull()
 class Interactive:
     def __init__(self):
         self.tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=HF_TOKEN_DOWNLOAD)
-        # 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).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()
-        # self.t = self.model.shared.weight[32097, 0].item()
+        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)
+        self.linear.bias = torch.nn.Parameter(self.model.shared.weight[32098, 0].unsqueeze(0)) # (1)
+        self.model.eval()
+        self.t = self.model.shared.weight[32097, 0].item()
 
     def run(self, statement):
-        # 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)
-        #     last_hidden_state = output.last_hidden_state.to(device) # (B=1, L, D)
-        #     hidden = last_hidden_state[0, -1, :] # (D)
-        #     logit = self.linear(hidden).squeeze(-1) # ()
-        #     logit_calibrated = logit / self.t
-        #     score = logit.sigmoid()
-        #     score_calibrated = logit_calibrated.sigmoid()
-        # return {
-        #     'logit': logit.item(),
-        #     'logit_calibrated': logit_calibrated.item(),
-        #     'score': score.item(),
-        #     'score_calibrated': score_calibrated.item(),
-        # }
+        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)
+            last_hidden_state = output.last_hidden_state.to(device) # (B=1, L, D)
+            hidden = last_hidden_state[0, -1, :] # (D)
+            logit = self.linear(hidden).squeeze(-1) # ()
+            logit_calibrated = logit / self.t
+            score = logit.sigmoid()
+            score_calibrated = logit_calibrated.sigmoid()
         return {
-            'logit': 0.0,
-            'logit_calibrated': 0.0,
-            'score': 0.5,
-            'score_calibrated': 0.5,
+            'logit': logit.item(),
+            'logit_calibrated': logit_calibrated.item(),
+            '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()