Upload llm.py

utils/llm.py

@@ -105,7 +105,8 @@ NGワード:
     next_token_logits = output.logits[0, -1, :]
     next_token_probs = torch.softmax(next_token_logits, dim=-1)
     sorted_ids = torch.argsort(next_token_probs, descending=True, dim=-1)
-  
+    
+    prob_d = {}
     for i in range(2):
       prob_d[tokenizer.decode(sorted_ids[i])] = next_token_probs[sorted_ids[i]].item()