0708_10

app.py

@@ -19,10 +19,10 @@ base_model = base_model
 tokenizer = AutoTokenizer.from_pretrained(base_model)
 
 sen_model = Classifier(base_model, num_labels=2, device='cpu', tokenizer=tokenizer)
-sen_model.load_state_dict(torch.load(sen_model_file))
+sen_model.load_state_dict(torch.load(sen_model_file, map_location=torch.device('cpu')))
 
 entity_model = Classifier(base_model, num_labels=2, device='cpu', tokenizer=tokenizer)
-entity_model.load_state_dict(torch.load(entity_model_file))
+entity_model.load_state_dict(torch.load(entity_model_file, map_location=torch.device('cpu')))
 
 
 def infer(test_sentence):
@@ -45,15 +45,15 @@ def infer(test_sentence):
         
         tokenized_data = tokenizer(form_, pair_, padding='max_length', max_length=512, truncation=True)
         
-        input_ids = torch.tensor([tokenized_data['input_ids']]).to(device)
-        attention_mask = torch.tensor([tokenized_data['attention_mask']]).to(device)
+        input_ids = torch.tensor([tokenized_data['input_ids']])
+        attention_mask = torch.tensor([tokenized_data['attention_mask']])
         
         first_sep = tokenized_data['input_ids'].index(2)
         last_sep = tokenized_data['input_ids'][first_sep+2:].index(2) + (first_sep + 2)        
         mask = [0] * len(tokenized_data['input_ids'])        
         for i in range(first_sep + 2, last_sep):
             mask[i] = 1     
-        mask = torch.tensor([mask]).to(device)
+        mask = torch.tensor([mask])
                 
         with torch.no_grad():
             outputs = entity_model(input_ids, attention_mask, mask)