Update optimum_neuron_export.py

optimum_neuron_export.py

@@ -134,17 +134,16 @@ def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discuss
             return discussion
     return None
 
-def export(model_id: str, task_or_pipeline:str, model_type: str, folder: str):
-
-    yield f"📦 Exporting model `{model_id}` for task `{task_or_pipeline}`..."
-
-    if model_type == "diffusers":
-        model_class = DIFFUSION_PIPELINE_MAPPING.get(task_or_pipeline)
-        
+def export(model_id: str, task_or_pipeline: str, model_type: str, folder: str):
+    """Export model to Neuron format. This is NOT a generator."""
+    
+    print(f"📦 Exporting model `{model_id}` for task `{task_or_pipeline}`...")
+    
     inputs = get_default_inputs(task_or_pipeline)
-    yield f"🔧 Using default inputs: {inputs}"
+    print(f"🔧 Using default inputs: {inputs}")
 
-    if task_or_pipeline in ENCODER_TASKS or SEQ2SEQ_TAKS:   
+    # ENCODER and SEQ2SEQ tasks
+    if task_or_pipeline in ENCODER_TASKS or task_or_pipeline in SEQ2SEQ_TAKS:   
         result = main_export(
             model_name_or_path=model_id,
             output=folder,
@@ -155,16 +154,22 @@ def export(model_id: str, task_or_pipeline:str, model_type: str, folder: str):
             **inputs,            
         )
 
-    if task_or_pipeline in DECODER_TASKS:
-        neuron_config = NeuronModelForCausalLM.get_neuron_config(model_name_or_path=model_id, **inputs)
+    # DECODER tasks
+    elif task_or_pipeline in DECODER_TASKS:
+        neuron_config = NeuronModelForCausalLM.get_neuron_config(
+            model_name_or_path=model_id, 
+            **inputs
+        )
         neuron_model = NeuronModelForCausalLM.export(
-            model_id=export_decoder_id, 
+            model_id=model_id,  # Fixed variable name
             neuron_config=neuron_config,
-            token = HF_TOKEN,
+            token=HF_TOKEN,
         )
-        model.save_pretrained(folder)
+        neuron_model.save_pretrained(folder)  # Fixed variable name
 
-    if task_or_pipeline in DIFFUSION_PIPELINE_MAPPING:
+    # DIFFUSION tasks
+    elif task_or_pipeline in DIFFUSION_PIPELINE_MAPPING:
+        model_class = DIFFUSION_PIPELINE_MAPPING.get(task_or_pipeline)
         model = model_class.from_pretrained(model_id)
         input_shapes = build_stable_diffusion_components_mandatory_shapes(**inputs)
         compiler_kwargs = {"auto_cast": "matmul", "auto_cast_type": "bf16"}
@@ -179,24 +184,49 @@ def export(model_id: str, task_or_pipeline:str, model_type: str, folder: str):
             model=model,
             **input_shapes,
         )
+    else:
+        raise ValueError(f"Unsupported task or pipeline: {task_or_pipeline}")
+    
+    print(f"✅ Export completed successfully to {folder}")
 
 
-def export_and_git_add(model_id: str, task_or_pipeline: str, model_type: str, folder: str, token: str) -> Any:
+def export_and_git_add(model_id: str, task_or_pipeline: str, model_type: str, folder: str, token: str):
+    """Export model and prepare git operations. This IS a generator."""
     
     try:
+        # Actually execute the export (not a generator anymore)
         export(model_id, task_or_pipeline, model_type, folder)
         yield "✅ Export completed successfully."
     except Exception as e:
         yield f"❌ Export failed with error: {e}"
         raise
 
+    # Verify that files were actually created
+    if not os.path.exists(folder) or not os.listdir(folder):
+        error_msg = f"❌ Export folder is empty or doesn't exist: {folder}"
+        yield error_msg
+        raise Exception(error_msg)
+    
+    yield f"📁 Found exported files in {folder}"
+
+    # Collect all files for git operations
     operations = []
+    file_count = 0
     for root, _, files in os.walk(folder):
         for filename in files:
             file_path = os.path.join(root, filename)
             repo_path = os.path.relpath(file_path, folder)
             operations.append(CommitOperationAdd(path_in_repo=repo_path, path_or_fileobj=file_path))
+            file_count += 1
+    
+    yield f"📦 Prepared {file_count} files for upload"
     
+    if file_count == 0:
+        error_msg = "❌ No files found to upload after export"
+        yield error_msg
+        raise Exception(error_msg)
+    
+    # Update model card
     try:
         card = ModelCard.load(model_id, token=token)
         if not hasattr(card.data, "tags") or card.data.tags is None:
@@ -213,10 +243,13 @@ def export_and_git_add(model_id: str, task_or_pipeline: str, model_type: str, fo
             readme_op.path_or_fileobj = readme_path
         else:
             operations.append(CommitOperationAdd(path_in_repo="README.md", path_or_fileobj=readme_path))
+        
+        yield "📝 Updated model card with neuron tag"
 
     except Exception as e:
         yield f"⚠️ Warning: Could not update model card: {e}"
 
+    # Return the operations
     yield ("__RETURN__", operations)
 
 def generate_neuron_repo_name(api, original_model_id: str, task_or_pipeline: str, token:str) -> str:
@@ -439,7 +472,7 @@ Generated using: [Optimum Neuron Compiler Space]({SPACES_URL})
         yield f"❌ Failed to create README PR: {e}"
         raise
 
-# --- Updated upload_to_custom_repo function (unchanged) ---
+# --- Updated upload_to_custom_repo function ---
 def upload_to_custom_repo(
     operations: List[CommitOperationAdd],
     custom_repo_id: str,