Update inference_fine_tune.py

inference_fine_tune.py

@@ -46,8 +46,8 @@ def generate_response(prompt:str):
        decoder_input = decoder_input.unsqueeze(0)
     temperature = 0.7
     top_k = 50
-
-    while decoder_input.shape[1] < 2000 :
+    i = 0
+    while i < 1024:
         # Apply causal mask based on current decoder_input length
         # decoder_mask = (decoder_input != pad_token_id).unsqueeze(0).int() & causal_mask(decoder_input.size(1)).type_as(input_mask).to(device)
         # Get model output
@@ -59,7 +59,7 @@ def generate_response(prompt:str):
         next_token = torch.multinomial(probs, num_samples=1)
         next_token = top_k_indices.gather(-1, next_token)
         word += tokenizer.decode([next_token.item()])
-        
+        i+=1
         decoder_input = torch.cat([decoder_input, next_token], dim=1)
         if decoder_input.shape[1] > config['seq_len']:
             decoder_input = decoder_input[:,-config['seq_len']:]