Upload inference.py

inference.py

@@ -0,0 +1,47 @@
+import sys
+import os
+import torch
+from transformers import AutoTokenizer
+from safetensors.torch import load_file
+
+model_dir = "/Users/Goekdeniz.Guelmez@computacenter.com/Library/CloudStorage/OneDrive-COMPUTACENTER/Desktop/mlx-lm/dev"
+sys.path.append(model_dir)
+
+from modeling_longcat_flash import LongcatFlashForCausalLM
+from configuration_longcat_flash import LongcatFlashConfig
+import json
+
+# Load model
+with open(os.path.join(model_dir, "config.json"), 'r') as f:
+    config_dict = json.load(f)
+
+config = LongcatFlashConfig(**config_dict)
+model = LongcatFlashForCausalLM(config)
+
+# Load weights
+state_dict = load_file(os.path.join(model_dir, "model.safetensors"))
+model.load_state_dict(state_dict)
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_dir)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# Test with "Lorem"
+text = "Lorem"
+inputs = tokenizer(text, return_tensors="pt", padding=True, return_attention_mask=True)
+
+print(f"Input: {text}")
+
+# Generate
+with torch.no_grad():
+    generated_ids = model.generate(
+        input_ids=inputs["input_ids"],
+        attention_mask=inputs["attention_mask"],
+        max_new_tokens=10,
+        do_sample=False,
+        pad_token_id=tokenizer.pad_token_id
+    )
+    generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+    
+print(f"Generated: {generated_text}")