Update app.py

app.py

@@ -9,8 +9,12 @@ project = os.environ["PROJECT"]
 wandb_api_key = os.environ["WANDB_API_KEY"]
 
 config = {
+    "max_output_tokens": 800,
     #"model": "text-bison@001",
     "model": "gemini-pro",
+    "temperature": 0,
+    "top_k": 250,
+    "top_p": 0.999,
 }
 
 wandb.login(key = wandb_api_key)
@@ -30,14 +34,14 @@ 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).text
+    completion = generation_model.predict(prompt = prompt).text
+    #completion = generation_model.generate_content(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."