Skip to content

Commit cecfec5

Browse files
huggingface: handle image-text-to-text pipeline task (#31611)
**Description:** Allows for HuggingFacePipeline to handle image-text-to-text pipeline
1 parent 50f998a commit cecfec5

File tree

1 file changed

+5
-2
lines changed

1 file changed

+5
-2
lines changed

libs/partners/huggingface/langchain_huggingface/llms/huggingface_pipeline.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
VALID_TASKS = (
2424
"text2text-generation",
2525
"text-generation",
26+
"image-text-to-text",
2627
"summarization",
2728
"translation",
2829
)
@@ -38,8 +39,8 @@ class HuggingFacePipeline(BaseLLM):
3839
3940
To use, you should have the ``transformers`` python package installed.
4041
41-
Only supports `text-generation`, `text2text-generation`, `summarization` and
42-
`translation` for now.
42+
Only supports `text-generation`, `text2text-generation`, `image-text-to-text`,
43+
`summarization` and `translation` for now.
4344
4445
Example using from_model_id:
4546
.. code-block:: python
@@ -327,6 +328,8 @@ def _generate(
327328
text = response["generated_text"]
328329
elif self.pipeline.task == "text2text-generation":
329330
text = response["generated_text"]
331+
elif self.pipeline.task == "image-text-to-text":
332+
text = response["generated_text"]
330333
elif self.pipeline.task == "summarization":
331334
text = response["summary_text"]
332335
elif self.pipeline.task in "translation":

0 commit comments

Comments
 (0)