Update app.py

app.py

@@ -10,8 +10,8 @@ wandb_api_key = os.environ["WANDB_API_KEY"]
 
 config = {
     "max_output_tokens": 800,
-    "model": "text-bison@001",
-    #"model": "gemini-pro",
+    #"model": "text-bison@001",
+    "model": "gemini-pro",
     "temperature": 0.1,
     "top_k": 40,
     "top_p": 1.0,
@@ -34,27 +34,27 @@ vertexai.init(project = project,
               credentials = credentials
              )
 
-from vertexai.language_models import TextGenerationModel
-generation_model = TextGenerationModel.from_pretrained(config.model)
-#from vertexai.preview.generative_models import GenerativeModel
-#generation_model = GenerativeModel(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(config.model)
 
 def invoke(prompt):
-    completion = generation_model.predict(prompt = prompt,
-                                          max_output_tokens = config.max_output_tokens,
-                                          temperature = config.temperature,
-                                          top_k = config.top_k,
-                                          top_p = config.top_p,
-                                         ).text
+    #completion = generation_model.predict(prompt = prompt,
+    #                                      max_output_tokens = config.max_output_tokens,
+    #                                      temperature = config.temperature,
+    #                                      top_k = config.top_k,
+    #                                      top_p = config.top_p,
+    #                                     ).text
     #completion = generation_model.generate_content(prompt, generation_config = {
     #                                                           "max_output_tokens": config.max_output_tokens,
     #                                                           "temperature": config.temperature,
     #                                                           "top_k": config.top_k,
     #                                                           "top_p": config.top_p,
     #                                                       }).text
-    wandb.log({"prompt": prompt, "completion": 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."
+    #wandb.log({"prompt": prompt, "completion": 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."
 
 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 gemini-pro foundation model. Model performance evaluation via <a href='https://wandb.ai/bstraehle'>Weights & Biases</a>."""