add back writing to llama

nod-ai · Nov 20, 2023 · 451d66c · 451d66c
1 parent 661fcb0
commit 451d66c
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/python/turbine_models/custom_models/stateless_llama.py b/python/turbine_models/custom_models/stateless_llama.py
@@ -324,11 +324,16 @@ def get_token_from_logits(logits):
     if args.run_vmfb:
         run_vmfb_comparison(args)
     else:
-        export_transformer_model(
+        mod_str, _ = export_transformer_model(
             args.hf_model_name,
             args.hf_auth_token,
             args.compile_to,
             args.external_weights,
             args.external_weight_file,
             args.quantization,
         )
+        safe_name = args.hf_model_name.split("/")[-1].strip()
+        safe_name = re.sub("-", "_", safe_name)
+        with open(f"{safe_name}.mlir", "w+") as f:
+            f.write(mod_str)
+