Added support for embedding models

Signed-off-by: amitraj <[email protected]>

Author: quic-amitraj

QEfficient/generation/text_generation_inference.py

@@ -367,7 +367,7 @@ def cloud_ai_100_exec_embed(
     return prefill_outputs
 
 
-class TextGeneration:
+class QEffTextGenerationBase:
     def __init__(
         self,
         tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],

QEfficient/transformers/models/modeling_auto.py

@@ -213,6 +213,13 @@ def export(self, export_dir: Optional[str] = None) -> str:
             example_inputs["batch_index"] = torch.arange(bs).view(bs, 1)
             dynamic_axes["batch_index"] = {0: "batch_size"}
 
+        return self._export(
+            example_inputs,
+            output_names,
+            dynamic_axes,
+            export_dir=export_dir,
+        )
+
     def compile(
         self,
         onnx_path: Optional[str] = None,
@@ -381,7 +388,7 @@ def generate(
         device_id: List[int] = [0],
         runtime_ai100: bool = True,
         seq_len: int = constants.Constants.CTX_LEN,
-    ):
+    ) -> str:
         if runtime_ai100:
             if not isinstance(self.qpc_path, Path):
                 raise TypeError("Please run compile API first!")