from mistral_common.tokens.tokenizers.mistral import MistralTokenizer 
from mistral_common.protocol.instruct.messages import UserMessage
from mistral_common.protocol.instruct.request import ChatCompletionRequest
from transformers import AutoModelForCausalLM
import torch


# Load Mistral tokenizer
model_name = "nemostral"
tokenizer = MistralTokenizer.from_model(model_name)

# Tokenize a list of messages
tokenized = tokenizer.encode_chat_completion(
    ChatCompletionRequest(
        messages=[
            UserMessage(content="How many peolpe live in France and all its neighbours? List all of them!")
        ],
        model=model_name,
    )
)
tokens, text = tokenized.tokens, tokenized.text

input_ids = torch.tensor([tokens]).to("cuda")

model = AutoModelForCausalLM.from_pretrained("./", torch_dtype=torch.bfloat16).to("cuda")


out = model.generate(input_ids, max_new_tokens=1024)

generated = out[0, input_ids.shape[-1]:-1].tolist()

print(tokenizer.decode(generated))