Update app.py

app.py

@@ -25,15 +25,15 @@ def generate_text(prompt, temperature, top_p):
     for _ in range(80):  # Adjust the range to control the number of tokens generated
         with torch.no_grad():
             outputs = model(input_tokens)
-            predictions = outputs.logits / temperature
-            sorted_logits, sorted_indices = torch.sort(predictions[:, -1, :], descending=True)
+            predictions = outputs.logits[:, -1, :] / temperature
+            sorted_logits, sorted_indices = torch.sort(predictions, descending=True)
             cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
             sorted_indices_to_remove = cumulative_probs > top_p
             sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
             sorted_indices_to_remove[..., 0] = 0
             indices_to_remove = sorted_indices[sorted_indices_to_remove]
-            predictions[:, -1, indices_to_remove] = -float('Inf')
-            next_token = torch.multinomial(torch.softmax(predictions[:, -1, :], dim=-1), 1)
+            predictions[:, indices_to_remove] = -float('Inf')
+            next_token = torch.multinomial(torch.softmax(predictions, dim=-1), 1)
         
         input_tokens = torch.cat((input_tokens, next_token), dim=1)