import os
import torch
from datetime import datetime, timedelta
from transformers import  Mistral3ForConditionalGeneration
from mistral_common.protocol.instruct.request import ChatCompletionRequest
from mistral_common.tokens.tokenizers.mistral import MistralTokenizer


def load_system_prompt(repo_id_or_path: str, filename: str) -> str:
    from huggingface_hub import hf_hub_download

    if os.path.isdir(repo_id_or_path):
        file_path = os.path.join(repo_id_or_path, filename)
    else:
        file_path = hf_hub_download(repo_id=repo_id_or_path, filename=filename)
    with open(file_path, "r") as file:
        system_prompt = file.read()
    today = datetime.today().strftime("%Y-%m-%d")
    yesterday = (datetime.today() - timedelta(days=1)).strftime("%Y-%m-%d")
    model_name = repo_id_or_path.split("/")[-1]
    return system_prompt.format(name=model_name, today=today, yesterday=yesterday)

model_id = "Intel/Magistral-Small-2509-int4-AutoRound"
model = Mistral3ForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
if os.path.isdir(model_id):
    tokenizer = MistralTokenizer.from_file(os.path.join(model_id, "tekken.json"))
else:
    tokenizer = MistralTokenizer.from_hf_hub(model_id)

SYSTEM_PROMPT = load_system_prompt(model.name_or_path, "SYSTEM_PROMPT.txt")

content = "Which of the depicted countries has the best food? Which the second and third and fourth? Name the country, its color on the map and one its city that is visible on the map, but is not the capital. Make absolutely sure to only name a city that can be seen on the map."
image_url = "https://huggingface.co/datasets/patrickvonplaten/random_img/resolve/main/europe.png"
conversation = [{"role": "system", "content": SYSTEM_PROMPT}]
conversation.append({
    "role": "user",
    "content": [{
        "type": "text",
        "text": content
    }, {
        "type": "image_url",
        "image_url": {
            "url": image_url
        }
    }],
},)
tokenized = tokenizer.encode_chat_completion(ChatCompletionRequest(messages=conversation, continue_final_message=False))
breakpoint()
input_ids = torch.tensor([tokenized.tokens])
attention_mask = torch.ones_like(input_ids)
pixel_values = torch.tensor(tokenized.images[0], dtype=torch.bfloat16).unsqueeze(0)
image_sizes = torch.tensor([pixel_values.shape[-2:]])
ret = {
    "input_ids": input_ids,
    "attention_mask": attention_mask,
    "pixel_values": pixel_values,
    "image_sizes": image_sizes,
}
for key in ret:
    if isinstance(ret[key], torch.Tensor):
        ret[key] = ret[key].to(model.device)
with torch.inference_mode():
    generation = model.generate(**ret, max_new_tokens=512, do_sample=True)

decoded = tokenizer.decode(generation)
print(decoded)


"""
Alright, the question is about identifying the best, second best, third best, and fourth best countries in terms of food from the depicted map. The user also wants the country's color on the map and a non-capital city visible on the map for each.

First, let's list out the countries visible on the map and their colors:

1. Iceland - light purple
2. Ireland - dark green
3. United Kingdom - yellow
4. France - brown
5. Spain - yellow
6. Portugal - yellow
7. Italy - orange
8. Greece - light blue
9. Germany - orange
10. Switzerland - light green
11. Austria - light green
12. Netherlands - orange
13. Belgium - light blue
14. Luxembourg - light blue
15. Denmark - light blue
16. Poland - light green
17. Czech Republic - purple
18. Slovakia - purple
19. Hungary - purple
20. Romania - light green
21. Bulgaria - dark brown
22. Russia - red
23. Finland - light green
24. Sweden - dark purple
25. Norway - dark green
26. Denmark - light blue
27. Estonia - light blue
28. Latvia - light blue
29. Lithuania - light blue
30. Belarus - dark purple
31. Ukraine - yellow
32. Moldova - yellow
33. Turkey - light blue (partially visible)

Now, let's consider the countries known for their food. This is subjective, but generally, Italy, France, Spain, and Greece are often praised for their cuisine.

1. Italy (orange) - Known for pasta, pizza, risotto, etc.
2. France (brown) - Known for wine, cheese, pastries, etc.
3. Spain (yellow) - Known for paella, tapas, jamón, etc.
4. Greece (light blue) - Known for gyros, souvlaki, feta, etc.

Now, we need to pick one non-capital city for each that is visible on the map.

1. Italy: Milan (not the capital, Rome is the capital)
2. France: Lyon (not the capital, Paris is the capital)
3. Spain: Barcelona (not the capital, Madrid is the capital)
4. Greece: Thessaloniki (not the capital, Athens is the capital)

Now, let's format the response:

1. First: Italy, color: orange, city: Milan
2. Second: France, color: brown, city: Lyon
3. Third: Spain, color: yellow, city: Barcelona
4. Fourth: Greece, color: light blue, city: Thessaloniki

But wait, the user asked for the best, second best, third best, and fourth best. We need to rank them. Let's assume the following ranking based on general popularity and global influence of cuisine:

1. Italy - often considered the best
2. France - often considered the second best
3. Spain - often considered the third best
4. Greece - often considered the fourth best

Now, let's prepare the response.1. **Italy** - Color: orange, City: Milan
2. **France** - Color: brown, City: Lyon
3. **Spain** - Color: yellow, City: Barcelona
4. **Greece** - Color: light blue, City: Thessaloniki

"""

Generate the model

https://github.com/intel/auto-round/pull/713 is required.

auto_round --model /models/Magistral-Small-2509/ --iters 200 --seqlen 1024 --output_dir tmp_autoround

Ethical Considerations and Limitations

The model can produce factually incorrect output, and should not be relied on to produce factually accurate information. Because of the limitations of the pretrained model and the finetuning datasets, it is possible that this model could generate lewd, biased or otherwise offensive outputs.

Therefore, before deploying any applications of the model, developers should perform safety testing.

Caveats and Recommendations

Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model.

Here are a couple of useful links to learn more about Intel's AI software:

Intel Neural Compressor link

Disclaimer

The license on this model does not constitute legal advice. We are not responsible for the actions of third parties who use this model. Please consult an attorney before using this model for commercial purposes.

Cite

@article{cheng2023optimize, title={Optimize weight rounding via signed gradient descent for the quantization of llms}, author={Cheng, Wenhua and Zhang, Weiwei and Shen, Haihao and Cai, Yiyang and He, Xin and Lv, Kaokao and Liu, Yi}, journal={arXiv preprint arXiv:2309.05516}, year={2023} }

arxiv github