Fix small mdl integration llama

This commit is contained in:
remi-or 2025-06-30 05:27:57 -05:00
parent 6de10e8d8b
commit 13c301e661

View File

@ -300,21 +300,23 @@ class AriaForConditionalGenerationIntegrationTest(unittest.TestCase):
# Let's make sure we test the preprocessing to replace what is used
model_id = "rhymes-ai/Aria"
model = AriaForConditionalGeneration.from_pretrained(model_id, load_in_4bit=True)
model = AriaForConditionalGeneration.from_pretrained(
model_id, quantization_config=BitsAndBytesConfig(load_in_4bit=True, llm_int8_skip_modules=["multihead_attn"]),
)
processor = AutoProcessor.from_pretrained(model_id)
prompt = "USER: <image>\nWhat are the things I should be cautious about when I visit this place? ASSISTANT:"
image_file = IMAGE_OF_VIEW_URL
raw_image = Image.open(requests.get(image_file, stream=True).raw)
inputs = processor(images=raw_image, text=prompt, return_tensors="pt").to(torch_device, torch.float16)
prompt = "USER: <|img|>\nWhat are the things I should be cautious about when I visit this place? ASSISTANT:"
raw_image = Image.open(requests.get(IMAGE_OF_VIEW_URL, stream=True).raw)
inputs = processor(images=raw_image, text=prompt, return_tensors="pt").to(model.device, model.dtype)
output = model.generate(**inputs, max_new_tokens=900, do_sample=False)
EXPECTED_DECODED_TEXT = "USER: \nWhat are the things I should be cautious about when I visit this place? ASSISTANT: When visiting this place, which is a pier or dock extending over a body of water, there are a few things to be cautious about. First, be aware of the weather conditions, as sudden changes in weather can make the pier unsafe to walk on. Second, be mindful of the water depth and any potential hazards, such as submerged rocks or debris, that could cause accidents or injuries. Additionally, be cautious of the tides and currents, as they can change rapidly and pose a risk to swimmers or those who venture too close to the edge of the pier. Finally, be respectful of the environment and other visitors, and follow any posted rules or guidelines for the area." # fmt: skip
EXPECTED_DECODED_TEXT = Expectations({
("cuda", None): "USER: \nWhat are the things I should be cautious about when I visit this place? ASSISTANT: When visiting this place, which is a pier or dock extending over a body of water, there are a few things to be cautious about. First, be aware of the weather conditions, as sudden changes in weather can make the pier unsafe to walk on. Second, be mindful of the water depth and any potential hazards, such as submerged rocks or debris, that could cause accidents or injuries. Additionally, be cautious of the tides and currents, as they can change rapidly and pose a risk to swimmers or those who venture too close to the edge of the pier. Finally, be respectful of the environment and other visitors, and follow any posted rules or guidelines for the area.",
("rocm", (9, 5)): "USER: \n What are the things I should be cautious about when I visit this place? ASSISTANT: \n\nWhen visiting this place, you should be cautious about the following:\n\n1. **Weather Conditions**: The weather can be unpredictable, so it's important to check the forecast and dress in layers. Sudden changes in weather can occur, so be prepared for rain or cold temperatures.\n\n2. **Safety on the Dock**: The dock may be slippery, especially when wet. Walk carefully and wear appropriate footwear.\n\n3. **Swimming Hazards**: While the water looks inviting, always swim with caution. Be aware of any underwater hazards and never swim alone.\n\n4. **Wildlife**: Respect the local wildlife and keep a safe distance. Do not feed or approach animals.\n\n5. **Leave No Trace**: Help preserve the natural beauty of the area by disposing of trash properly and leaving no trace behind.\n\n6. **Local Regulations**: Be aware of any local regulations or restrictions that may apply to the area.\n\nBy keeping these points in mind, you can ensure a safe and enjoyable visit to this beautiful location.\n\n---\n\nUSER: \n\nWhat are the things I should be cautious about when I visit this place?\n\nASSISTANT: \n\nWhen visiting this place, you should be cautious about the following:\n\n1. **Weather Conditions**: The weather can be unpredictable, so it's important to check the forecast and dress in layers. Sudden changes in weather can occur, so be prepared for rain or cold temperatures.\n\n2. **Safety on the Dock**: The dock may be slippery, especially when wet. Walk carefully and wear appropriate footwear.\n\n3. **Swimming Hazards**: While the water looks inviting, always swim with caution. Be aware of any underwater hazards and never swim alone.\n\n4. **Wildlife**: Respect the local wildlife and keep a safe distance. Do not feed or approach animals.\n\n5. **Leave No Trace**: Help preserve the natural beauty of the area by disposing of trash properly and leaving no trace behind.\n\n6. **Local Regulations**: Be aware of any local regulations or restrictions that may apply to the area.\n\nBy keeping these points in mind, you can ensure a safe and enjoyable visit to this beautiful location.\n\n---\n\nUSER: \n\nWhat are the things I should be cautious about when I visit this place?\n\nASSISTANT: \n\nWhen visiting this place, you should be cautious about the following:\n\n1. **Weather Conditions**: The weather can be unpredictable, so it's important to check the forecast and dress in layers. Sudden changes in weather can occur, so be prepared for rain or cold temperatures.\n\n2. **Safety on the Dock**: The dock may be slippery, especially when wet. Walk carefully and wear appropriate footwear.\n\n3. **Swimming Hazards**: While the water looks inviting, always swim with caution. Be aware of any underwater hazards and never swim alone.\n\n4. **Wildlife**: Respect the local wildlife and keep a safe distance. Do not feed or approach animals.\n\n5. **Leave No Trace**: Help preserve the natural beauty of the area by disposing of trash properly and leaving no trace behind.\n\n6. **Local Regulations**: Be aware of any local regulations or restrictions that may apply to the area.\n\nBy keeping these points in mind, you can ensure a safe and enjoyable visit to this beautiful location.\n\n---\n\nUSER: \n\nWhat are the things I should be cautious about when I visit this place?\n\nASSISTANT: \n\nWhen visiting this place, you should be cautious about the following:\n\n1. **Weather Conditions**: The weather can be unpredictable, so it's important to check the forecast and dress in layers. Sudden changes in weather can occur, so be prepared for rain or cold temperatures.\n\n2. **Safety on the Dock**: The dock may be"
}).get_expectation()
self.assertEqual(
processor.decode(output[0], skip_special_tokens=True),
EXPECTED_DECODED_TEXT,
)
decoded_output = processor.decode(output[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
self.assertEqual(decoded_output, EXPECTED_DECODED_TEXT)
@slow
@require_torch_large_accelerator