added the option to choose the model

autodistill · Jun 12, 2024 · cb1da92 · cb1da92
1 parent 9e361e0
commit cb1da92
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -51,6 +51,7 @@ base_model = Gemini(
     ),
     gcp_region="us-central1",
     gcp_project="project-name",
+    model="gemini-1.5-flash"
 )
 
 # run inference on an image

diff --git a/autodistill_gemini/gemini_model.py b/autodistill_gemini/gemini_model.py
@@ -13,18 +13,26 @@
 
 @dataclass
 class Gemini(DetectionBaseModel):
+    AVAILABLE_MODELS = ["gemini-1.5-flash", "gemini-1.5-pro", "gemini-pro-vision"]
     ontology: CaptionOntology
     api_key: str
     gcp_region: str
     gcp_project: str
+    model: str
 
     def __init__(
-        self, ontology: CaptionOntology, gcp_region: str, gcp_project: str
+        self, ontology: CaptionOntology, gcp_region: str, gcp_project: str, model: str
     ) -> None:
         self.ontology = ontology
         self.gcp_region = gcp_region
         self.gcp_project = gcp_project
 
+        if model in self.AVAILABLE_MODELS:
+            self.model = model
+        else:
+            raise ValueError(f"Choose one of the available models from {available_models}")
+
+
     def predict(
         self, input: str, prompt: str = "", confidence: int = 0.5
     ) -> sv.Detections:
@@ -40,7 +48,7 @@ def predict(
 
         vertexai.init(project=self.gcp_project, location=self.gcp_region)
 
-        multimodal_model = GenerativeModel("gemini-pro-vision")
+        multimodal_model = GenerativeModel(self.model)
 
         response = multimodal_model.generate_content(
             [prompt, Image.load_from_file(input)]