Spaces:
Sleeping
Sleeping
Ashmi Banerjee
commited on
Commit
·
c526665
1
Parent(s):
38f3034
models update
Browse files- 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.
|
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 |
-
|
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
|