fixed docstrings and unused imports

Author: elizjo

ads/aqua/extension/recommend_handler.py

@@ -1,11 +1,9 @@
-
 from tornado.web import HTTPError
 
 from ads.aqua.common.decorator import handle_exceptions
 from ads.aqua.extension.base_handler import AquaAPIhandler
 from ads.aqua.extension.errors import Errors
 from ads.aqua.shaperecommend.recommend import AquaRecommendApp
-from ads.config import COMPARTMENT_OCID
 
 
 class AquaRecommendHandler(AquaAPIhandler):
@@ -14,8 +12,6 @@ class AquaRecommendHandler(AquaAPIhandler):
 
     Methods
     -------
-    get(self, id: Union[str, List[str]])
-        Retrieves a list of AQUA deployments or model info or logs by ID.
     post(self, *args, **kwargs)
         Obtains the eligible compute shapes that would fit the specifed model, context length, model weights, and quantization level.
 
@@ -27,16 +23,15 @@ class AquaRecommendHandler(AquaAPIhandler):
     @handle_exceptions
     def post(self, *args, **kwargs):  # noqa: ARG002
         """
-        Lists the eligible GPU compute shapes for the specifed model.
+        Obtains the eligible compute shapes that would fit the specifed model, context length, model weights, and quantization level.
 
         Returns
         -------
-        List[ComputeShapeSummary]:
-            The list of the model deployment shapes.
+        ShapeRecommendationReport
+            Report containing shape recommendations and troubleshooting advice, if any.
         """
         try:
             input_data = self.get_json_body()
-            # input_data["compartment_id"] = self.get_argument("compartment_id", default=COMPARTMENT_OCID)
         except Exception as ex:
             raise HTTPError(400, Errors.INVALID_INPUT_DATA_FORMAT) from ex
 
@@ -45,6 +40,7 @@ def post(self, *args, **kwargs):  # noqa: ARG002
 
         self.finish(AquaRecommendApp().which_gpu(**input_data))
 
+
 __handlers__ = [
     ("recommendation/?([^/]*)", AquaRecommendHandler),
 ]

ads/aqua/shaperecommend/constants.py

@@ -37,12 +37,8 @@
 TEXT_GENERATION = "text_generation"
 SAFETENSORS = "safetensors"
 
-#TODO:
-SHAPES_METADATA = "/Users/elizjo/tmp/accelerated-data-science/ads/aqua/resources/gpu_shapes_index.json"
-
 TROUBLESHOOT_MSG = "The selected model is too large to fit on standard GPU shapes with the current configuration.\nAs troubleshooting, we have suggested the two largest available GPU shapes using the smallest quantization level ('4bit') to maximize chances of fitting the model. "
 
-TEXT_MODEL = "text-generation"
 
 QUANT_MAPPING = {
             "float32": 4,

ads/aqua/shaperecommend/recommend.py

@@ -18,7 +18,6 @@
 )
 from ads.aqua.shaperecommend.constants import (
     SAFETENSORS,
-    SHAPES_METADATA,
     TEXT_GENERATION,
     TROUBLESHOOT_MSG,
 )
@@ -126,6 +125,9 @@ def get_model_config(ocid: str):
         AquaValueError
             If the OCID is not for a Data Science model, or if the model type is not supported,
             or if required files/tags are not present.
+
+        AquaRecommendationError
+            If the model OCID provided is not supported (only text-generation decoder models in safetensor format supported).
         """
         resource_type = get_resource_type(ocid)
 
@@ -176,9 +178,7 @@ def get_model_config(ocid: str):
         return data
 
     @staticmethod
-    def valid_compute_shapes(
-        file: str = SHAPES_METADATA,
-    ) -> List["ComputeShapeSummary"]:
+    def valid_compute_shapes() -> List["ComputeShapeSummary"]:
         """
         Returns a filtered list of GPU-only ComputeShapeSummary objects by reading and parsing a JSON file.
 
@@ -306,7 +306,7 @@ def summarize_shapes_for_seq_lens(
 
         troubleshoot_msg = ""
 
-        if len(recommendations) > 5:
+        if len(recommendations) > 2:
             recommendations = ShapeReport.pareto_front(recommendations)
 
         if not recommendations: