Fix bug in getting vocab_size and missing ccl in forward

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

Author: quic-xiyushi

QEfficient/transformers/models/modeling_auto.py

@@ -8,7 +8,7 @@
 import warnings
 from pathlib import Path
 from time import perf_counter
-from typing import Dict, List, Optional, Union
+from typing import List, Optional, Union
 
 import numpy as np
 import torch
@@ -1081,7 +1081,7 @@ def export(
                 output_names=output_names["lang"],
                 dynamic_axes=dynamic_axes["lang"],
                 continuous_batching=self.continuous_batching,
-                vocab_size=self.config.vocab_size,
+                vocab_size=self.model.language_model.config.vocab_size,
                 qaic_config=self.lang_model.model.qaic_config,
             )
 

QEfficient/transformers/sampler/sampler.py

@@ -105,6 +105,7 @@ def sampler_forward(
     attention_mask: Optional[torch.Tensor] = None,
     position_ids: Optional[torch.LongTensor] = None,
     past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
+    comp_ctx_lengths: Optional[torch.LongTensor] = None,
     batch_index: Optional[torch.LongTensor] = None,
     inputs_embeds: Optional[torch.FloatTensor] = None,
     labels: Optional[torch.LongTensor] = None,
@@ -181,6 +182,7 @@ def sampler_forward(
             position_ids=position_ids,
             image_idx=image_idx,
             past_key_values=past_key_values,
+            comp_ctx_lengths=comp_ctx_lengths,
         )
         if batch_index is not None:
             forward_kwargs["batch_index"] = batch_index
@@ -195,6 +197,7 @@ def sampler_forward(
             attention_mask=attention_mask,
             position_ids=position_ids,
             past_key_values=past_key_values,
+            comp_ctx_lengths=comp_ctx_lengths,
             batch_index=batch_index,
             inputs_embeds=inputs_embeds,
             use_cache=use_cache,

tests/transformers/sampler/test_sampler.py

@@ -6,10 +6,10 @@
 # -----------------------------------------------------------------------------
 
 from typing import List, Union
-from transformers import AutoConfig, AutoProcessor
 
 import numpy as np
 import pytest
+from transformers import AutoProcessor
 
 from QEfficient import QEFFAutoModelForCausalLM, QEFFAutoModelForImageTextToText
 from QEfficient.generation.cloud_infer import QAICInferenceSession