github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-05-31 20:19:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

Feature/enhance huggingfacepipeline to handle different return type (#11394)

Browse Source 
**Description:** Avoid huggingfacepipeline to truncate the response if
user setup return_full_text as False within huggingface pipeline.

**Dependencies:** : None
**Tag maintainer:**   Maybe @sam-h-bean ?

---------

Co-authored-by: Bagatur <baskaryan@gmail.com>
This commit is contained in:
hsuyuming 2023-10-11 20:09:03 -06:00 committed by GitHub
parent 2aba9ab47e
commit 0b743f005b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 17 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
libs/langchain/langchain/llms/huggingface_pipeline.py
View File

@ -202,8 +202,23 @@ class HuggingFacePipeline(BaseLLM):
response = response[0]
if self.pipeline.task == "text-generation":
# Text generation return includes the starter text
text = response["generated_text"][len(batch_prompts[j]) :]
try:
from transformers.pipelines.text_generation import ReturnType
remove_prompt = (
self.pipeline._postprocess_params.get("return_type")
!= ReturnType.NEW_TEXT
)
except Exception as e:
logger.warning(
f"Unable to extract pipeline return_type. "
f"Received error:\n\n{e}"
)
remove_prompt = True
if remove_prompt:
text = response["generated_text"][len(batch_prompts[j]) :]
else:
text = response["generated_text"]
elif self.pipeline.task == "text2text-generation":
text = response["generated_text"]
elif self.pipeline.task == "summarization":