address PR comments

Author: CalebCourier

guardrails/llm_providers.py

@@ -487,6 +487,7 @@ def get_llm_ask(llm_api: Callable, *args, **kwargs) -> PromptCallableBase:
     try:
         from transformers import (  # noqa: F401 # type: ignore
             FlaxPreTrainedModel,
+            GenerationMixin,
             PreTrainedModel,
             TFPreTrainedModel,
         )
@@ -498,14 +499,24 @@ def get_llm_ask(llm_api: Callable, *args, **kwargs) -> PromptCallableBase:
             or isinstance(api_self, TFPreTrainedModel)
             or isinstance(api_self, FlaxPreTrainedModel)
         ):
-            return HuggingFaceModelCallable(*args, model_generate=llm_api, **kwargs)
+            if (
+                hasattr(llm_api, "__func__")
+                and llm_api.__func__ == GenerationMixin.generate
+            ):
+                return HuggingFaceModelCallable(*args, model_generate=llm_api, **kwargs)
+            raise ValueError("Only text generation models are supported at this time.")
     except ImportError:
         pass
     try:
         from transformers import Pipeline  # noqa: F401 # type: ignore
 
         if isinstance(llm_api, Pipeline):
-            return HuggingFacePipelineCallable(*args, pipeline=llm_api, **kwargs)
+            # Couldn't find a constant for this
+            if llm_api.task == "text-generation":
+                return HuggingFacePipelineCallable(*args, pipeline=llm_api, **kwargs)
+            raise ValueError(
+                "Only text generation pipelines are supported at this time."
+            )
     except ImportError:
         pass
 

tests/unit_tests/test_llm_providers.py

@@ -650,7 +650,7 @@ def test_get_llm_ask_anthropic():
     not importlib.util.find_spec("transformers"),
     reason="transformers is not installed",
 )
-def test_get_llm_ask_hugging_face_model():
+def test_get_llm_ask_hugging_face_model(mocker):
     from transformers import PreTrainedModel
 
     from guardrails.llm_providers import HuggingFaceModelCallable
@@ -661,9 +661,6 @@ class MockModel(PreTrainedModel):
         def __init__(self, *args, **kwargs):
             self._modules = {}
 
-        def generate(self, *args, **kwargs):
-            pass
-
     mock_model = MockModel()
 
     prompt_callable = get_llm_ask(mock_model.generate)
@@ -681,6 +678,8 @@ def test_get_llm_ask_hugging_face_pipeline():
     from guardrails.llm_providers import HuggingFacePipelineCallable
 
     class MockPipeline(Pipeline):
+        task = "text-generation"
+
         def __init__(self, *args, **kwargs):
             pass