--- library_name: transformers tags: [] --- ``` model: opt-125m config: Float8DynamicActivationFloat8WeightConfig config version: 1 torchao version: 0.13.dev ``` ``` import torch import io from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig from huggingface_hub import HfApi model_id = "facebook/opt-125m" from torchao.quantization import Float8DynamicActivationFloat8WeightConfig, PerRow quant_config = Float8DynamicActivationFloat8WeightConfig(granularity=PerRow(), version=1) quantization_config = TorchAoConfig(quant_type=quant_config) quantized_model = AutoModelForCausalLM.from_pretrained( model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config, ) tokenizer = AutoTokenizer.from_pretrained(model_id) # Push to hub USER_ID = "torchao-testing" MODEL_NAME = model_id.split("/")[-1] save_to = f"{USER_ID}/{MODEL_NAME}-Float8DynamicActivationFloat8WeightConfig-v1-0.13.dev" quantized_model.push_to_hub(save_to, safe_serialization=False) tokenizer.push_to_hub(save_to) # Manual Testing prompt = "Hey, are you conscious? Can you talk to me?" print("Prompt:", prompt) inputs = tokenizer( prompt, return_tensors="pt", ).to("cuda") # setting temperature to 0 to make sure result deterministic generated_ids = quantized_model.generate(**inputs, max_new_tokens=128, temperature=0) api = HfApi() buf = io.BytesIO() torch.save(prompt, buf) api.upload_file( path_or_fileobj=buf, path_in_repo="model_prompt.pt", repo_id=save_to, ) buf = io.BytesIO() torch.save(generated_ids, buf) api.upload_file( path_or_fileobj=buf, path_in_repo="model_output.pt", repo_id=save_to, ) output_text = tokenizer.batch_decode( generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False ) print("Response:", output_text[0][len(prompt) :]) ```