Ashmi Banerjee commited on
Commit
c526665
·
1 Parent(s): 38f3034

models update

Browse files
Files changed (1) hide show
  1. models/gemini.py +3 -7
models/gemini.py CHANGED
@@ -3,17 +3,13 @@ import sys
3
 
4
  sys.path.append("../")
5
  from setup.vertex_ai_setup import initialize_vertexai_params
6
- from vertexai.preview.generative_models import GenerativeModel
7
 
8
 
9
  def get_gemini_response(prompt_text, model, parameters: Optional = None) -> str:
10
  initialize_vertexai_params()
11
  if model is None or parameters is None:
12
  model = "gemini-1.0-pro"
13
- parameters = {
14
- "max_output_tokens": 1024,
15
- "temperature": 0.2
16
- }
17
- model = GenerativeModel.from_pretrained(model, **parameters)
18
  model_response = model.generate_content(prompt_text)
19
- return model_response
 
3
 
4
  sys.path.append("../")
5
  from setup.vertex_ai_setup import initialize_vertexai_params
6
+ from vertexai.generative_models import GenerativeModel
7
 
8
 
9
  def get_gemini_response(prompt_text, model, parameters: Optional = None) -> str:
10
  initialize_vertexai_params()
11
  if model is None or parameters is None:
12
  model = "gemini-1.0-pro"
13
+ model = GenerativeModel(model)
 
 
 
 
14
  model_response = model.generate_content(prompt_text)
15
+ return model_response.text