Update llava.py (#54)

added stopping_criteria from gen kwargs
IntelLabs · Jul 14, 2024 · 36b7ecb · 36b7ecb
1 parent 722860d
commit 36b7ecb
Showing 1 changed file with 4 additions and 0 deletions.
diff --git 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 = "<image>"