Update README.md

README.md

@@ -13,6 +13,8 @@ import logging
 
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
+from huggingface_hub import HfApi
+import io
 
 # Configure logging to see warnings and debug information
 logging.basicConfig(
@@ -47,41 +49,41 @@ tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 # Push to hub
 MODEL_NAME = model_id.split("/")[-1]
-save_to = f"torchao-testing/{MODEL_NAME}-IntxWeightOnlyConfig-v{version}-0.14.0.dev-safetensors"
+save_to = f"torchao-testing/{MODEL_NAME}-IntxWeightOnlyConfig-v{version}-0.14.0.dev"
 quantized_model.push_to_hub(save_to, safe_serialization=False)
 tokenizer.push_to_hub(save_to)
 
 
 # Manual Testing
-prompt = "What are we having for dinner?"
+prompt = "Hey, are you conscious? Can you talk to me?"
 print("Prompt:", prompt)
 inputs = tokenizer(
     prompt,
     return_tensors="pt",
 ).to("cuda")
 
-# Detting temperature to 0 to make sure result deterministic
+# setting temperature to 0 to make sure result deterministic
 generated_ids = quantized_model.generate(**inputs, max_new_tokens=128, temperature=0)
 
-correct_output_text = tokenizer.batch_decode(
-    generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
+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,
 )
-print("Response:", correct_output_text[0][len(prompt) :])
-
 
-# Load model from saved checkpoint
-reloaded_model = AutoModelForCausalLM.from_pretrained(
-    save_to,
-    device_map="auto",
-    torch_dtype=torch.bfloat16,
+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,
 )
 
-generated_ids = reloaded_model.generate(**inputs, max_new_tokens=128, temperature=0)
 output_text = tokenizer.batch_decode(
     generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
 )
 print("Response:", output_text[0][len(prompt) :])
-
-assert(correct_output_text == output_text)
-
 ```