Update app.py

app.py

@@ -34,19 +34,24 @@ 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).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
+    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."