fix

app.py

@@ -7,7 +7,7 @@ from transformers import GPT2LMHeadModel, GPT2Tokenizer
 
 device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
 tokenizer = GPT2Tokenizer.from_pretrained("dendee-geco_test-on-zuco1.0_gpt2_tmptoken_TRT_bs32_lr1e6_linearLR")
-model = GPT2LMHeadModel.from_pretrained('dendee-geco_test-on-zuco1.0_gpt2_tmptoken_TRT_bs32_lr1e6_linearLR', return_dict=True)
+model = GPT2LMHeadModel.from_pretrained("dendee-geco_test-on-zuco1.0_gpt2_tmptoken_TRT_bs32_lr1e6_linearLR", return_dict=True)
 model.to(device)
 
 
@@ -21,9 +21,10 @@ def calculate_surprisals(
     ]
     input_ids = tokenizer.encode(input_text, return_tensors='pt').to(device)
 
-    logits = model(input_ids)['logits'].squeeze(0)
+    logits = model(input_ids)['logits'].squeeze(0) # (1, seq_len)
+    logprob = torch.log_softmax(logits, dim=-1)
     # can't calculate surprisals for the first token, hence 0
-    surprisals = [0] + (- torch.gather(logits[:-1, :], -1, input_ids[:, 1:]).squeeze(0)).tolist()
+    surprisals = [0] + (- torch.gather(logprob[:-1, :], -1, input_ids[:, 1:]).squeeze(0)).tolist()
     mean_surprisal = np.mean(surprisals[1:])
 
     if normalize_surprisals:
@@ -53,7 +54,7 @@ def highlight_token(token: str, score: float):
 def create_highlighted_text(tokens2scores: List[Tuple[str, float]]):
     highlighted_text: str = ""
     for token, score in tokens2scores:
-        highlighted_text += highlight_token(token, score)
+        highlighted_text += highlight_token(token, score) + ' '
     highlighted_text += "<br><br>"
     return highlighted_text