Update optimum_neuron_export.py

optimum_neuron_export.py

@@ -149,8 +149,10 @@ DIFFUSION_PIPELINE_MAPPING = {
 
 def get_default_inputs(task_or_pipeline: str, pipeline_name: str = None) -> Dict[str, int]:
     """Get default input shapes based on task type or diffusion pipeline type."""
-    if task_or_pipeline in ["feature-extraction", "sentence-transformers", "fill-mask", "question-answering", "text-classification", "token-classification","text-generation"]:
+    if task_or_pipeline in ["feature-extraction", "sentence-transformers", "fill-mask", "question-answering", "text-classification", "token-classification"]:
         return {"batch_size": 1, "sequence_length": 128}
+    elif task_or_pipeline == "text_generation":
+        return {"batch_size": 1, "sequence_length": 128, "tensor_parallel_size": 4}
     elif task_or_pipeline == "multiple-choice":
         return {"batch_size": 1, "num_choices": 4, "sequence_length": 128}
     elif task_or_pipeline == "text2text-generation":
@@ -278,9 +280,7 @@ def export_transformer_model(model_id: str, task: str, folder: str, token: str)
     inputs = get_default_inputs(task)
     compiler_configs = {"auto_cast": "matmul", "auto_cast_type": "bf16", "instance_type": "inf2"}
     yield f"🔧 Using default inputs: {inputs}"
-    
-    # Clear any old cache artifacts before export
-     
+         
     try:
         # Trigger the export/compilation
         model = model_class.from_pretrained(
@@ -315,21 +315,60 @@ def export_transformer_model(model_id: str, task: str, folder: str, token: str)
         yield f"❌ Export failed with error: {e}"
         raise
 
+def export_decoder_model(model_id: str, folder: str, token:str) -> Generator: 
+    """Export decoder-only models using optimum.neuron"""
+    
+    yield f"📦 Exporting decoder model `{model_id}` ..."
+    
+    export_kwargs = get_default_inputs("text_generation")
+    yield f"🔧 Using default export_kwargs: {export_kwargs}"
+         
+    try:
+        # Trigger the export/compilation
+        neuron_config = NeuronModelForCausalLM.get_neuron_config(model_name_or_path=model_id, **export_kwargs)
+        neuron_config.target = "inf2"
+        model = NeuronModelForCausalLM.export(
+                model_id=model_id, neuron_config=neuron_config,
+            )
+        
+        yield "✅ Export/compilation completed successfully."
+        
+        # Find the newly created cache artifacts
+        yield "🔍 Locating compiled artifacts in Neuron cache..."
+        cache_artifact_dir = find_neuron_cache_artifacts(cache_base_dir)
+        
+        if not cache_artifact_dir:
+            raise Exception("❌ Could not find compiled artifacts in Neuron cache")
+        
+        yield f"📂 Found artifacts at: {cache_artifact_dir}"
+        
+        # Copy artifacts from cache to our target folder
+        yield f"📋 Copying artifacts to export folder..."
+        if os.path.exists(folder):
+            shutil.rmtree(folder)
+        shutil.copytree(cache_artifact_dir, folder)
+        
+        yield f"✅ Artifacts successfully copied to {folder}"
+        
+    except Exception as e:
+        yield f"❌ Export failed with error: {e}"
+        raise
+    
 def export_and_git_add(model_id: str, task_or_pipeline: str, model_type: str, folder: str, token: str, pipeline_name: str = None) -> Any:
 
     operations = []
     
     try:
         if model_type == "diffusers":
-            # For diffusion models, use the new export function
-            export_gen = export_diffusion_model(model_id, pipeline_name, task_or_pipeline, folder, token)
-            for message in export_gen:
-                yield message
+            export_gen = export_diffusion_model(model_id, pipeline_name, task_or_pipeline, folder, token) 
         else:
-            # For transformer models, use the existing function
-            export_gen = export_transformer_model(model_id, task_or_pipeline, folder, token)
-            for message in export_gen:
-                yield message
+            if task_or_pipeline == "text_generation":
+                export_gen = export_decoder_model(model_id, folder, token)
+            else:   
+                export_gen = export_transformer_model(model_id, task_or_pipeline, folder, token)
+        
+        for message in export_gen:
+            yield message
         
         # Create operations from exported files
         for root, _, files in os.walk(folder):