models update

models/gemini.py

@@ -3,17 +3,13 @@ import sys
 
 sys.path.append("../")
 from setup.vertex_ai_setup import initialize_vertexai_params
-from vertexai.preview.generative_models import GenerativeModel
+from vertexai.generative_models import GenerativeModel
 
 
 def get_gemini_response(prompt_text, model, parameters: Optional = None) -> str:
     initialize_vertexai_params()
     if model is None or parameters is None:
         model = "gemini-1.0-pro"
-        parameters = {
-            "max_output_tokens": 1024,
-            "temperature": 0.2
-        }
-    model = GenerativeModel.from_pretrained(model, **parameters)
+    model = GenerativeModel(model)
     model_response = model.generate_content(prompt_text)
-    return model_response
+    return model_response.text