Add unit tests

Signed-off-by: quic-xiyushi <xiyushi@qti.qualcomm.com>

Author: quic-xiyushi

QEfficient/generation/vlm_generation.py

@@ -36,6 +36,7 @@
     write_io_files,
 )
 from QEfficient.utils import LRUCache
+from QEfficient.utils.constants import Constants
 from QEfficient.utils.logging_utils import logger
 
 
@@ -303,6 +304,13 @@ def _execute_chunked_prefill(
             prefill_ccl_id = 0
             lang_inputs["comp_ctx_lengths"] = self.list_of_comp_ctx_lengths_prefill[prefill_ccl_id]
 
+        if self.include_sampler:
+            for op in Constants.SAMPLER_OPS:
+                if decode_batch_id is not None:
+                    lang_inputs[op] = self.sampling_params[op][decode_batch_id.flatten()]
+                else:
+                    lang_inputs[op] = self.sampling_params[op]
+
         for i in range(num_chunks):
             input_ids_slice = lang_inputs["input_ids"][:, i * self._prefill_seq_len : (i + 1) * self._prefill_seq_len]
             position_ids_slice = lang_inputs["position_ids"][
@@ -328,6 +336,11 @@ def _execute_chunked_prefill(
 
                 chunk_inputs["comp_ctx_lengths"] = lang_inputs["comp_ctx_lengths"]
 
+            if self.include_sampler:
+                chunk_inputs["last_accepted_output_tokens"] = chunk_inputs["input_ids"]
+                for op in Constants.SAMPLER_OPS:
+                    chunk_inputs[op] = lang_inputs[op]
+
             outputs = self._session.run(chunk_inputs)
 
             if "image_idx_output" in outputs:

QEfficient/transformers/models/modeling_auto.py

@@ -881,7 +881,10 @@ def __init__(
             If `full_batch_size` is provided.
         """
         if kwargs.pop("full_batch_size", None):
-            raise NotImplementedError("Continuous batching is not supported for image-text-to-text models yet.")
+            continuous_batching = True
+            warnings.warn(
+                "full_batch_size argument is deprecated. Use continuous_batching=True instead.", DeprecationWarning, 2
+            )
         self.model = model
         self.config = model.config
 
@@ -1028,7 +1031,7 @@ def export(
                 output_names=output_names["lang"],
                 dynamic_axes=dynamic_axes["lang"],
                 continuous_batching=self.continuous_batching,
-                vocab_size=self.lang_model.model.config.vocab_size,
+                vocab_size=self.config.vocab_size,
                 qaic_config=self.lang_model.model.qaic_config,
             )
 
@@ -1235,6 +1238,7 @@ def generate(
         device_ids: List[int] = None,
         runtime_ai100: bool = True,
         generation_len: Optional[int] = None,
+        **kwargs,
     ) -> Union[torch.Tensor, np.ndarray]:
         """
         Generates output by executing the compiled QPC(s) on Cloud AI 100 Hardware cards.
@@ -1293,6 +1297,7 @@ def generate(
                 full_batch_size=fbs,
                 comp_ctx_lengths_prefill=self.comp_ctx_lengths_prefill,
                 comp_ctx_lengths_decode=self.comp_ctx_lengths_decode,
+                **kwargs,
             )
 
             # Call generate method
@@ -1572,11 +1577,16 @@ def __init__(
         Raises
         ------
         NotImplementedError
-            If `full_batch_size` is provided.
+            If `full_batch_size` is provided or `continuous_batching` is True or `include_sampler` is True.
         """
         if kwargs.pop("full_batch_size", None):
+            warnings.warn(
+                "full_batch_size argument is deprecated. Use continuous_batching=True instead.", DeprecationWarning, 2
+            )
+            raise NotImplementedError("Continuous batching is not supported for image-text-to-text models yet.")
+        if kwargs.pop("continuous_batching", None):
             raise NotImplementedError("Continuous batching is not supported for image-text-to-text models yet.")
-        if kwargs.pop("qaic_config", None):
+        if qaic_config is not None and qaic_config.pop("include_sampler", False):
             raise NotImplementedError("On-device sampling is not supported for single QPC multimodal models yet.")
         super().__init__(model, **kwargs)