Update README.md

README.md

@@ -66,24 +66,15 @@ CUDA_VISIBLE_DEVICES=0,1 python eval.py \
 ## Usage
 
 ```python
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from peft import PeftModel, PeftConfig
-
-peft_model_id = 'SeanLee97/angle-llama-7b-nli-20231027'
-config = PeftConfig.from_pretrained(peft_model_id)
-tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
-model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path).bfloat16().cuda()
-model = PeftModel.from_pretrained(model, peft_model_id).cuda()
-
-def decorate_text(text: str):
-    return f'Summarize sentence "{text}" in one word:"'
-
-inputs = 'hello world!'
-tok = tokenizer([decorate_text(inputs)], return_tensors='pt')
-for k, v in tok.items():
-    tok[k] = v.cuda()
-vec = model(output_hidden_states=True, **tok).hidden_states[-1][:, -1].float().detach().cpu().numpy()
+from angle_emb import AnglE
+
+angle = AnglE.from_pretrained('NousResearch/Llama-2-7b-hf', pretrained_lora_path='SeanLee97/angle-llama-7b-nli-20231027')
+angle.set_prompt()
+print('prompt:', angle.prompt)
+vec = angle.encode({'text': 'hello world'}, to_numpy=True)
 print(vec)
+vecs = angle.encode([{'text': 'hello world1'}, {'text': 'hello world2'}], to_numpy=True)
+print(vecs)
 ```
 
 ## Citation