Update app.py

app.py

@@ -18,7 +18,7 @@ def generate_text(prompt):
     input_tokens = tokenizer.encode(prompt, return_tensors='pt')
     input_tokens = input_tokens.to('cpu')
     
-    generated_tokens = []
+    generated_text = prompt  # Start with the initial prompt
 
     for _ in range(80):  # Adjust the range to control the number of tokens generated
         with torch.no_grad():
@@ -26,15 +26,11 @@ def generate_text(prompt):
             predictions = outputs.logits
             next_token = torch.multinomial(torch.softmax(predictions[:, -1, :], dim=-1), 1)
         
-        generated_tokens.append(next_token.item())
         input_tokens = torch.cat((input_tokens, next_token), dim=1)
         
         decoded_token = tokenizer.decode(next_token.item())
-        yield decoded_token  # Yield each token as it is generated
-    
-    # Decode the entire generated text
-    generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
-    yield generated_text
+        generated_text += decoded_token  # Append the new token to the generated text
+        yield generated_text  # Yield the entire generated text so far
 
 # Create a Gradio interface with a text input and a text output
 interface = gr.Interface(fn=generate_text, inputs='text', outputs='text', live=True)