From 36b7ecbe8497f0801780fa738e47fac9f10b6642 Mon Sep 17 00:00:00 2001 From: Moshe Berchansky Date: Sun, 14 Jul 2024 11:16:02 +0300 Subject: [PATCH] Update llava.py (#54) added stopping_criteria from gen kwargs --- fastrag/generators/llava.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/fastrag/generators/llava.py b/fastrag/generators/llava.py index 600de31..c4aad5f 100644 --- a/fastrag/generators/llava.py +++ b/fastrag/generators/llava.py @@ -105,6 +105,10 @@ def __init__( if "streamer" not in self.generation_kwargs: self.generation_kwargs["streamer"] = TextStreamer(AutoTokenizer.from_pretrained(model)) + if "stopping_criteria" in self.generation_kwargs: + self.stopping_criteria_list = self.generation_kwargs["stopping_criteria"] + del self.generation_kwargs["stopping_criteria"] + self.processor = AutoProcessor.from_pretrained(model) self.image_token = ""