Update inference_fine_tune.py

inference_fine_tune.py

@@ -47,7 +47,7 @@ def generate_response(prompt:str):
     temperature = 0.7
     top_k = 50
     i = 0
-    while i < 1024:
+    while decoder_input.shape[1] < 2000:
         # 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
@@ -63,7 +63,7 @@ def generate_response(prompt:str):
         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']:]
-        if next_token.item() == eos_token_id:
+        if next_token.item() == eos_token_id  or i >= 1024:
             break
     print("Output : ",word)
     return word