Update app.py

app.py

@@ -33,7 +33,9 @@ from vertexai.language_models import TextGenerationModel
 generation_model = TextGenerationModel.from_pretrained("text-bison@001")
 
 def invoke(prompt):
-    return generation_model.predict(prompt = prompt).text
+    completion = generation_model.predict(prompt = prompt).text
+    wandb.log({"prompt": prompt, "completion": completion})
+    return completion
 
 gr.close_all()
 demo = gr.Interface(fn=invoke,