yujunjun/haystack
1
0
Fork 0
mirror of https://github.com/deepset-ai/haystack.git synced 2025-12-02 09:56:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: Allow to set num_beams in HFInvocationLayer (#4731)

Browse Source 
Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>
Co-authored-by: bogdankostic <bogdankostic@web.de>
This commit is contained in:
Wang, Yi 2023-04-25 22:08:06 +08:00 committed by GitHub
parent 7fa3591f5f
commit 2be1a68fce
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
haystack/nodes/prompt/invocation_layer/hugging_face.py
View File

@ -193,7 +193,10 @@ class HFLocalInvocationLayer(PromptModelInvocationLayer):
model_input_kwargs["stopping_criteria"] = StoppingCriteriaList([sw])
if top_k:
model_input_kwargs["num_return_sequences"] = top_k
model_input_kwargs["num_beams"] = top_k
if "num_beams" not in model_input_kwargs or model_input_kwargs["num_beams"] < top_k:
if "num_beams" in model_input_kwargs:
logger.warning("num_beams should not be less than top_k, hence setting it to %s", top_k)
model_input_kwargs["num_beams"] = top_k
# max_new_tokens is used for text-generation and max_length for text2text-generation
if is_text_generation:
model_input_kwargs["max_new_tokens"] = self.max_length