Update app.py

app.py

@@ -11,7 +11,7 @@ HF_TOKEN = os.environ.get("HF_TOKEN")
 tokenizer = GPT2Tokenizer.from_pretrained('gpt2', use_auth_token=HF_TOKEN)
 model = GPT2LMHeadModel.from_pretrained('skylersterling/TopicGPT', use_auth_token=HF_TOKEN)
 model.eval()
-model.to("cpu")
+model.to('cpu')
 
 # Define the function that generates text from a prompt
 def generate_text(prompt):
@@ -30,13 +30,12 @@ def generate_text(prompt):
         input_tokens = torch.cat((input_tokens, next_token), dim=1)
         
         decoded_token = tokenizer.decode(next_token.item())
-        # Print each token as it is generated
-        print(decoded_token, end='', flush=True)
+        yield decoded_token  # Yield each token as it is generated
     
-    # Decode the generated tokens to a string
+    # Decode the entire generated text
     generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
-    return generated_text
+    yield generated_text
 
 # Create a Gradio interface with a text input and a text output
-interface = gr.Interface(fn=generate_text, inputs='text', outputs='text')
+interface = gr.Interface(fn=generate_text, inputs='text', outputs='text', live=True)
 interface.launch()