Update app.py

app.py

@@ -17,9 +17,6 @@ config = {
     "top_p": 1.0,
 }
 
-#wandb.login(key = wandb_api_key)
-#wandb.init(project = "vertex-ai-llm", config = config)
-
 credentials = json.loads(credentials)
 
 from google.oauth2 import service_account
@@ -38,6 +35,12 @@ vertexai.init(project = project,
 from vertexai.preview.generative_models import GenerativeModel
 generation_model = GenerativeModel(config["model"])
 
+def wandb_log(prompt, completion):
+    wandb.login(key = wandb_api_key)
+    wandb.init(project = "vertex-ai-llm", config = config)
+    wandb.log({"prompt": prompt, "completion": completion})
+    wandb.finish()
+
 def invoke(prompt):
     #completion = generation_model.predict(prompt = prompt,
     #                                      max_output_tokens = config["max_output_tokens"],
@@ -51,7 +54,7 @@ def invoke(prompt):
     #                                                           "top_k": config["top_k"],
     #                                                           "top_p": config["top_p"],
     #                                                       }).text
-    #wandb.log({"prompt": prompt, "completion": completion})
+    #wandb_log(prompt, completion)
     #return completion
     return "🛑 Execution is commented out. To view the source code see https://huggingface.co/spaces/bstraehle/google-vertex-ai-llm/tree/main."