Update README.md

README.md

@@ -20,16 +20,20 @@ License: [MIT License](https://opensource.org/license/mit/)
 
 How to use
 ```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
 base_model = 'llama-2-7b'
-comp_method = 'sparsegpt_unstructured'
 comp_degree = 0.1
-model_path = f'vita-group/comp-{base_model}_{comp_method}_s{comp_degree}'
+comp_method = 'sparsegpt_unstructured'
+model_path = f'vita-group/comp-{arch}_{comp_method}_s{comp_degree}'
 model = AutoModelForCausalLM.from_pretrained(
         model_path, 
         torch_dtype=torch.float16, 
         low_cpu_mem_usage=True, 
         device_map="auto"
     )
+tokenizer = AutoTokenizer.from_pretrained('meta-llama/Llama-2-7b')
+input_ids = tokenizer('Hello! I am a VITA-compressed-LLM chatbot!', return_tensors='pt').input_ids
+outputs = model.generate(input_ids)
 ```