Please conver the following code to work with rocm. import torch
import os
from transformers import AutoModelForCausalLM, AutoTokenizer, TextGenerationPipeline

# --- Configuration ---
MODEL_PATH = r"G:\model\Promt-generator"  # Your local model directory
PROMPT = "test story"
OUTPUT_TEXT_NAME = "test1.txt"
DTYPE = torch.float16  # Try float16 first, fallback to float32 if needed

# --- Check if model path exists ---
if not os.path.isdir(MODEL_PATH):
    print(f"Error: Model path not found: {MODEL_PATH}")
    print("Please ensure the path is correct and contains the model files.")
    exit()

# --- Device Setup (Attempt DirectML, fallback to CPU) ---
device = None
try:
    import torch_directml
    if torch_directml.is_available():
        device = torch_directml.device()
        print(f"Successfully initialized DirectML. Using device: {device}")
    else:
        print("DirectML is installed, but no suitable device was found or it's not enabled.")
        print("Falling back to CPU. This will be very slow.")
        device = torch.device("cpu")
except ImportError:
    print("torch-directml is not installed.")
    print("Please install it by running: pip install torch-directml")
    print("Falling back to CPU. This will be very slow.")
    device = torch.device("cpu")
except Exception as e:
    print(f"An error occurred while trying to initialize DirectML: {e}")
    print("Falling back to CPU. This will be very slow.")
    device = torch.device("cpu")

print(f"Selected device: {device}")

# --- 1. Load the model and tokenizer ---
print(f"Loading model from: {MODEL_PATH}")
try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, local_files_only=True)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        torch_dtype=DTYPE,
        local_files_only=True,
    )
    print("Model and tokenizer loaded successfully.")
except Exception as e:
    print(f"Error loading the model: {e}")
    if "fp16" in str(e).lower() and DTYPE == torch.float16:
        print("Try changing DTYPE to torch.float32 if your GPU or model has issues with float16.")
    exit()

# --- 2. Send to the selected device (DirectML or CPU) ---
print(f"Sending model to device: {device}...")
try:
    model = model.to(device)
    print("Model sent to device successfully.")
except Exception as e:
    print(f"Error sending model to device: {e}")
    if DTYPE == torch.float16 and ("half" in str(e).lower() or "fp16" in str(e).lower()):
        print("The model or parts of it might not support float16 on DirectML with your current setup.")
        print("Try changing DTYPE to torch.float32 in the script and re-run.")
    exit()

# --- 3. Generate text ---
print(f"Generating text with prompt: \"{PROMPT}\"")
try:
    pipe = TextGenerationPipeline(model=model, tokenizer=tokenizer, device=0 if device.type != "cpu" else -1)
    # You can adjust max_new_tokens, temperature, etc. as needed
    result = pipe(PROMPT, max_new_tokens=100, do_sample=True, temperature=0.8)
    generated_text = result[0]['generated_text']
    print("Text generated successfully.")
except Exception as e:
    print(f"Error during text generation: {e}")
    if "implemented" in str(e).lower() and DTYPE == torch.float16:
         print("An operation might not be implemented for float16 on DirectML.")
         print("Try changing DTYPE to torch.float32 in the script and re-run.")
    exit()

# --- 4. Save the result ---
try:
    with open(OUTPUT_TEXT_NAME, "w", encoding="utf-8") as f:
        f.write(generated_text)
    print(f"Text saved successfully as {OUTPUT_TEXT_NAME}!")
except Exception as e:
    print(f"Error saving text: {e}")