import torch
from transformers import AutoTokenizer
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel

# Load the model and tokenizer
model_name = "Trelis/mamba-2.8b-slimpj-chat-4k"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = MambaLMHeadModel.from_pretrained(model_name, dtype=torch.bfloat16, device="cuda")

# Define the prompt
prompt = "what languages do you speak? answer me in english"

# Initialize an empty list for messages
messages = []

# Append the prompt to the messages list as a dictionary
messages.append(dict(role="user", content=prompt))

device='cuda'

formatted=tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
print(formatted)

# Prepare the input for the model
input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True).to(device)
max_length = input_ids.shape[1] + 500  # Assuming you want to generate 100 tokens

# Generate function for Mamba model
def generate_mamba(input_ids, max_length):
    return model.generate(
        input_ids=input_ids,
        max_length=max_length,
        cg=True,
        return_dict_in_generate=True,
        output_scores=True,
        enable_timing=False,
        temperature=0.01,
        top_k=1,
        top_p=1.0,
        eos_token_id=tokenizer.eos_token_id,
        repetition_penalty=1.0,
    )

# Run the generation
out = generate_mamba(input_ids, max_length)

# Decode and print the generated text
decoded_sequences = tokenizer.batch_decode(out.sequences.tolist())
for sequence in decoded_sequences:
    generated_text = sequence[len(tokenizer.decode(input_ids[0], skip_special_tokens=False)):]
    print(generated_text)