Update app.py

app.py

@@ -10,7 +10,6 @@ wandb_api_key = os.environ["WANDB_API_KEY"]
 
 config = {
     "max_output_tokens": 800,
-    #"model": "text-bison@001",
     "model": "gemini-pro",
     "temperature": 0.1,
     "top_k": 40,
@@ -30,8 +29,6 @@ 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"])
 
@@ -44,39 +41,38 @@ def wandb_log(prompt, completion):
 def invoke(prompt):
     if (prompt == ""):
         raise gr.Error("Prompt is required.")
+
     completion = ""
+    
     try:
-        #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"],
-        #                                     )
-        #if (completion.text != None):
-        #    completion = completion.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"],
-                                                               })
+        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"],
+                                                       })
+        
         if (completion.text != None):
             completion = completion.text
     except Exception as e:
         completion = e
+        
         raise gr.Error(e)
     finally:
         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."
 
 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. RAG evaluation via <a href='https://wandb.ai/bstraehle'>Weights & Biases</a>."""
 
 gr.close_all()
-demo = gr.Interface(fn=invoke, 
+
+demo = gr.Interface(fn = invoke, 
                     inputs = [gr.Textbox(label = "Prompt", lines = 1)],
                     outputs = [gr.Textbox(label = "Completion", lines = 1)],
                     title = "Generative AI - LLM",
                     description = description)
+
 demo.launch()