Refactor generative model initialization to use gemini-1.5-flash

kkdai · Jul 7, 2024 · ebc103c · ebc103c
1 parent a28b7ea
commit ebc103c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/main.py b/main.py
@@ -188,7 +188,7 @@ def generate_gemini_text_complete(prompt):
     """
     Generate a text completion using the generative model.
     """
-    model = genai.GenerativeModel('gemini-pro')
+    model = genai.GenerativeModel('gemini-1.5-flash')
     response = model.generate_content(prompt)
     return response
 
@@ -201,7 +201,7 @@ def generate_json_from_receipt_image(img, prompt):
     :param prompt: prompt for the generative model.
     :return: the generated JSON representation of the receipt data.
     """
-    model = genai.GenerativeModel('gemini-pro-vision')
+    model = genai.GenerativeModel('gemini-1.5-flash')
     response = model.generate_content([prompt, img], stream=True)
     response.resolve()
     return response