2424For multimodal models, use --image-paths argument to provide image path(s),\
2525\
2626
27-
27+ For multimodal logits, since the model is not suppose to generate image, the logits correspond to images \
28+
29+
2830More examples:
2931python3 -m MaxText.scratch_code.generate_hf_golden_logits --model-id=meta-llama/Llama-4-Scout-17B-16E \
3032\
31- =true --gcs-bucket=<bucket> --hf-model-path=<hf_checkpoint_path> \
32-
33+ \
34+ --output-format=pickle
3335
3436python3 -m MaxText.scratch_code.generate_hf_golden_logits --model-id=google/gemma-3-4b-it \
3537\
36- apply-chat-template=false -- gcs-bucket=<bucket> --hf-model-path=<hf_checkpoint_path> \
38+ \
3739
3840"""
3941
4042import torch
4143import argparse
42- from transformers import AutoTokenizer , AutoProcessor , AutoModelForCausalLM
44+ from transformers import AutoTokenizer , AutoProcessor
4345import jsonlines
46+ import pickle
47+ import numpy as np
4448from google .cloud import storage
4549from PIL import Image
4650
@@ -55,78 +59,84 @@ def upload_blob(bucket_name, source_file_name, destination_blob_name):
5559 blob .upload_from_filename (source_file_name )
5660
5761
58- def save_golden_logits (model_id , output_path , prompt_texts , apply_chat_template , gcs_bucket , hf_model_path , image_paths ):
62+ def save_golden_logits (
63+ model_id , output_path , prompt_texts , apply_chat_template , gcs_bucket , hf_model_path , image_paths , output_format
64+ ):
5965 """save golden logits"""
6066 if hf_model_path is None :
6167 hf_model_path = model_id
68+
69+ if model_id .startswith ("meta-llama/Llama-4" ):
70+ from transformers import Llama4ForConditionalGeneration # pylint: disable=import-outside-toplevel
71+
72+ model_class = Llama4ForConditionalGeneration
73+ else :
74+ from transformers import AutoModelForCausalLM # pylint: disable=import-outside-toplevel
75+
76+ model_class = AutoModelForCausalLM
77+
6278 tokenizer = AutoTokenizer .from_pretrained (model_id )
63- model = AutoModelForCausalLM .from_pretrained (
79+ print (f"loading model from { hf_model_path } )
80+ model = model_class .from_pretrained (
6481 hf_model_path ,
6582 torch_dtype = torch .float32 ,
6683 trust_remote_code = True ,
6784 )
6885
6986 all_data_to_save = []
7087 for i , prompt_text in enumerate (prompt_texts ):
71- # Encode the prompt text
88+ # 1. Prepare inputs for the model and base data for saving
89+ data_to_save = {"prompt" : prompt_text }
7290 if image_paths :
73- try :
74- image = Image .open (image_paths [i ])
75- except Exception as e :
76- raise e
77- image = image .convert ("RGB" )
78- # TODO (aireenmei): remove this when Llama-4 supports dynamic image shapes.
91+ image = Image .open (image_paths [i ]).convert ("RGB" )
7992 if model_id .startswith ("meta-llama/Llama-4" ):
8093 image = image .resize ((336 , 336 ))
81- processor = AutoProcessor .from_pretrained (model_id , token = True )
94+ processor = AutoProcessor .from_pretrained (model_id , token = True ) if image_paths else None
8295 if apply_chat_template :
83- messages = [
84- {
85- "role" : "user" ,
86- "content" : [
87- {"type" : "image" },
88- {"type" : "text" , "text" : prompt_text },
89- ],
90- },
91- ]
96+ messages = [{"role" : "user" , "content" : [{"type" : "image" }, {"type" : "text" , "text" : prompt_text }]}]
9297 formatted_prompt = processor .apply_chat_template (messages , tokenize = False , add_generation_prompt = True )
9398 inputs = processor (text = formatted_prompt , images = image , return_tensors = "pt" )
9499 else :
95100 formatted_prompt = prompt_text
96101 inputs = processor (text = formatted_prompt , images = image , return_tensors = "pt" , add_special_tokens = False )
97- with torch .no_grad ():
98- outputs = model (** inputs )
99- logits = outputs .logits .cpu ().numpy ().astype ("float32" )
100-
101- data_to_save = {
102- "prompt" : prompt_text ,
103- "formatted_prompt" : formatted_prompt ,
104- "tokens" : inputs ["input_ids" ].tolist ()[0 ],
105- "attention_mask" : inputs ["attention_mask" ].tolist ()[0 ],
106- "image_path" : image_paths [i ],
107- "pixel_values" : inputs ["pixel_values" ].tolist ()[0 ],
108- "logits" : logits .tolist ()[0 ],
109- }
102+
103+ inputs ["pixel_values" ] = inputs ["pixel_values" ].to (torch .float32 )
104+ print (f"{ apply_chat_template = } { formatted_prompt = } )
105+ data_to_save .update ({"formatted_prompt" : formatted_prompt , "image_path" : image_paths [i ]})
110106 else :
111107 input_ids = tokenizer .encode (prompt_text , return_tensors = "pt" )
112- # Get the logits for the prompt + completion
113- with torch .no_grad ():
114- outputs = model (input_ids )
115- logits = outputs .logits .cpu ().numpy ().astype ("float32" )
116-
117- # Prepare data to be saved
118- data_to_save = {
119- "prompt" : prompt_text ,
120- "tokens" : input_ids .tolist ()[0 ],
121- "logits" : logits .tolist ()[0 ], # Convert numpy array to list for JSON serialization
122- }
108+ inputs = {"input_ids" : input_ids }
109+
110+ # 2. Run inference
111+ with torch .no_grad ():
112+ outputs = model (** inputs )
113+ logits = outputs .logits .cpu ().numpy ().astype ("float32" )
114+
115+ # 3. Populate final data dictionary with tensors from inputs and logits
116+ for key , value in inputs .items ():
117+ new_key = "tokens" if key == "input_ids" else key
118+ data_to_save [new_key ] = value .cpu ().numpy ()[0 ]
119+ data_to_save ["logits" ] = logits [0 ]
120+
123121 print (f"Token length is { len (data_to_save ['tokens' ])} { prompt_text } )
124122 print (f"raw ids: { data_to_save ['tokens' ]} )
123+
124+ # 4. Convert numpy arrays to lists if format is json
125+ if output_format == "json" :
126+ for key , value in data_to_save .items ():
127+ if isinstance (value , np .ndarray ):
128+ data_to_save [key ] = value .tolist ()
129+
125130 all_data_to_save .append (data_to_save )
126131
127- with jsonlines .open (output_path , "w" ) as f :
128- f .write_all (all_data_to_save )
129- print (f"File is stored locally at { output_path } )
132+ # 5. Save the collected data
133+ if output_format == "json" :
134+ with jsonlines .open (output_path , "w" ) as f :
135+ f .write_all (all_data_to_save )
136+ elif output_format == "pickle" :
137+ with open (output_path , "wb" ) as f :
138+ pickle .dump (all_data_to_save , f )
139+ print (f"File is stored locally at { output_path } )
130140
131141 if gcs_bucket :
132142 upload_blob (gcs_bucket , output_path , f"golden-logits/{ model_id } { output_path } )
@@ -140,10 +150,8 @@ def main(raw_args=None) -> None:
140150 parser .add_argument ("--prompts" , type = str , required = True , help = "A semicolon-separated list of prompts." )
141151 parser .add_argument (
142152 "--apply-chat-template" ,
143- type = bool ,
144- required = False ,
145- default = False ,
146- help = "Whether to apply chat template from the HF processor. Used for image+text input." ,
153+ action = "store_true" ,
154+ help = "Apply chat template from the HF processor. Use for image+text input. Pass this flag to enable; omit to disable." ,
147155 )
148156 parser .add_argument (
149157 "--gcs-bucket" , type = str , required = False , default = None , help = "A GCS bucket to store logits, without gs://."
@@ -154,6 +162,13 @@ def main(raw_args=None) -> None:
154162 parser .add_argument (
155163 "--image-paths" , type = str , required = False , default = None , help = "A semicolon-separated list of image_paths."
156164 )
165+ parser .add_argument (
166+ "--output-format" ,
167+ type = str ,
168+ required = False ,
169+ default = "json" ,
170+ help = "The output format for the golden logits. (json, pickle)" ,
171+ )
157172 args = parser .parse_args (raw_args )
158173 prompts = args .prompts .split (";" )
159174 image_paths = args .image_paths .split (";" ) if args .image_paths else []
@@ -164,7 +179,14 @@ def main(raw_args=None) -> None:
164179 if args .apply_chat_template :
165180 assert image_paths , "apply_chat_template is only used for image+text input, so image_paths must be provided."
166181 save_golden_logits (
167- args .model_id , args .output_path , prompts , args .apply_chat_template , args .gcs_bucket , args .hf_model_path , image_paths
182+ args .model_id ,
183+ args .output_path ,
184+ prompts ,
185+ args .apply_chat_template ,
186+ args .gcs_bucket ,
187+ args .hf_model_path ,
188+ image_paths ,
189+ args .output_format ,
168190 )
169191
170192
0 commit comments