Update generate.py

generate.py

@@ -24,13 +24,14 @@ def generate_response(prompt, max_length=100, top_k=40):
 
             # Top-k sampling
             top_k_probs, top_k_indices = torch.topk(next_token_logits, top_k)
-            probs = torch.softmax(top_k_probs.squeeze(0), dim=-1)  # Flatten
+            probs = torch.softmax(top_k_probs.squeeze(0), dim=-1)
             sampled_index = torch.multinomial(probs, 1).item()
             next_token = top_k_indices[0, sampled_index]
 
-        input_ids = torch.cat([input_ids, next_token.unsqueeze(0)], dim=1)
+        # Reshape next_token to match input_ids shape
+        next_token = next_token.unsqueeze(0).unsqueeze(0)  # Shape: (1, 1)
+        input_ids = torch.cat([input_ids, next_token], dim=1)
 
-        # Stop if EOS token
         if next_token.item() == tokenizer.eos_token_id:
             break