Update app.py

app.py

@@ -50,32 +50,33 @@ class GPT(nn.Module):
         return logits
 
     def generate(self, input_ids, max_new_tokens, temperature, top_k):
-    # Implement the text generation logic
-    output_ids = input_ids
-    for _ in range(max_new_tokens):
-        logits = self.forward(output_ids[:, -1:])
-        logits = logits / temperature
-        probs = F.softmax(logits, dim=-1)
-        
-        # Ensure probs is 2D
-        if probs.dim() == 3:
-            probs = probs.squeeze(0)  # Remove the batch dimension if it exists
-        
-        top_k_probs, top_k_indices = torch.topk(probs, k=top_k)
-        
-        # Ensure top_k_probs is 2D
-        if top_k_probs.dim() == 1:
-            top_k_probs = top_k_probs.unsqueeze(0)
-        
-        next_token = torch.multinomial(top_k_probs, num_samples=1)
-        next_token = top_k_indices.gather(-1, next_token)
-        
-        # Ensure next_token is 2D
-        if next_token.dim() == 1:
-            next_token = next_token.unsqueeze(0)
-        
-        output_ids = torch.cat([output_ids, next_token], dim=1)
-    return output_ids
+        # Implement the text generation logic
+        output_ids = input_ids
+        for _ in range(max_new_tokens):
+            logits = self.forward(output_ids[:, -1:])
+            logits = logits / temperature
+            probs = F.softmax(logits, dim=-1)
+            
+            # Ensure probs is 2D
+            if probs.dim() == 3:
+                probs = probs.squeeze(0)  # Remove the batch dimension if it exists
+            
+            top_k_probs, top_k_indices = torch.topk(probs, k=top_k)
+            
+            # Ensure top_k_probs is 2D
+            if top_k_probs.dim() == 1:
+                top_k_probs = top_k_probs.unsqueeze(0)
+            
+            next_token = torch.multinomial(top_k_probs, num_samples=1)
+            next_token = top_k_indices.gather(-1, next_token)
+            
+            # Ensure next_token is 2D
+            if next_token.dim() == 1:
+                next_token = next_token.unsqueeze(0)
+            
+            output_ids = torch.cat([output_ids, next_token], dim=1)
+        return output_ids
+
 # Initialize global variables
 model = None
 tokenizer = None