Update app.py

app.py

@@ -8,9 +8,9 @@ credentials = os.environ["CREDENTIALS"]
 project = os.environ["PROJECT"]
 wandb_api_key = os.environ["WANDB_API_KEY"]
 
-config = {
-    "model": "text-bison@001",
-}
+#config = {
+#    "model": "text-bison@001",
+#}
 
 wandb.login(key = wandb_api_key)
 wandb.init(project = "vertex-ai-llm", config = config)
@@ -29,14 +29,18 @@ vertexai.init(project = project,
               credentials = credentials
              )
 
-from vertexai.language_models import TextGenerationModel
-generation_model = TextGenerationModel.from_pretrained(config.model)
+#from vertexai.language_models import TextGenerationModel
+#generation_model = TextGenerationModel.from_pretrained(config.model)
+
+from vertexai.preview.generative_models import GenerativeModel
+generation_model = GenerativeModel("gemini-pro")
 
 def invoke(prompt):
     #completion = generation_model.predict(prompt = prompt).text
-    #wandb.log({"prompt": prompt, "completion": completion})
-    #return completion
-    return "Execution is commented out, but you can view the code at https://huggingface.co/spaces/bstraehle/google-vertex-ai-llm/tree/main."
+    completion = generation_model.generate_content(prompt = prompt).text
+    wandb.log({"prompt": prompt, "completion": completion})
+    return completion
+    #return "Execution is commented out, but you can view the code at https://huggingface.co/spaces/bstraehle/google-vertex-ai-llm/tree/main."
 
 description = """<a href='https://www.gradio.app/'>Gradio</a> UI using <a href='https://cloud.google.com/vertex-ai?hl=en/'>Google Vertex AI</a> API 
                  with Bison foundation model. Model performance evaluation via <a href='https://wandb.ai/bstraehle'>Weights & Biases</a>."""