From c87882c97922f85f9afdaa155f57e5e0081d2486 Mon Sep 17 00:00:00 2001
From: alex-matton <alexandre@cohere.com>
Date: Fri, 4 Aug 2023 11:27:34 +0200
Subject: [PATCH] expose basemodel

---
 cohere/responses/custom_model.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/cohere/responses/custom_model.py b/cohere/responses/custom_model.py
index c04fb9220..f4d41876a 100644
--- a/cohere/responses/custom_model.py
+++ b/cohere/responses/custom_model.py
@@ -82,6 +82,7 @@ def __init__(
         model_type: CUSTOM_MODEL_TYPE,
         created_at: datetime,
         completed_at: Optional[datetime],
+        base_model: Optional[str] = None,
         model_id: Optional[str] = None,
         hyperparameters: Optional[HyperParameters] = None,
     ) -> None:
@@ -92,6 +93,7 @@ def __init__(
         self.model_type = model_type
         self.created_at = created_at
         self.completed_at = completed_at
+        self.base_model = base_model
         self.model_id = model_id
         self.hyperparameters = hyperparameters
 
@@ -104,6 +106,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "CustomModel":
             model_type=REVERSE_CUSTOM_MODEL_PRODUCT_MAPPING[data["settings"]["finetuneType"]],
             created_at=_parse_date(data["created_at"]),
             completed_at=_parse_date(data["completed_at"]) if "completed_at" in data else None,
+            base_model=data["settings"]["baseModel"],
             model_id=data["model"]["route"] if "model" in data else None,
             hyperparameters=HyperParameters.from_response(data["settings"]["hyperparameters"])
             if data["settings"]["hyperparameters"]