Update app.py

app.py

@@ -59,20 +59,20 @@ with open("model.pt", "wb") as f:
     f.write(response.content)
 f.close()
 
-# model_load_path = "model.pt"
-# model = BertForSTS()
-# model.load_state_dict(torch.load(model_load_path))
-# model.to(device)
-
-# def predict_similarity(sentence_pair):
-#   test_input = tokenizer(sentence_pair, padding='max_length', max_length = param_max_length, truncation=True, return_tensors="pt").to(device)
-#   test_input['input_ids'] = test_input['input_ids']
-#   test_input['attention_mask'] = test_input['attention_mask']
-#   del test_input['token_type_ids']
-#   output = model(test_input)
-#   sim = torch.nn.functional.cosine_similarity(output[0], output[1], dim=0).item()*2-1
-
-#   return sim
+model_load_path = "model.pt"
+model = BertForSTS()
+model.load_state_dict(torch.load(model_load_path))
+model.to(device)
+
+def predict_similarity(sentence_pair):
+  test_input = tokenizer(sentence_pair, padding='max_length', max_length = param_max_length, truncation=True, return_tensors="pt").to(device)
+  test_input['input_ids'] = test_input['input_ids']
+  test_input['attention_mask'] = test_input['attention_mask']
+  del test_input['token_type_ids']
+  output = model(test_input)
+  sim = torch.nn.functional.cosine_similarity(output[0], output[1], dim=0).item()*2-1
+
+  return sim
 
 # Create a Gradio interface with a text input zone
 iface = gr.Interface(