Update README.md

README.md

@@ -42,8 +42,8 @@ pip install transformers==4.51.0
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+# Load model and tokenizer
 model_name = "infly/inf-query-aligner"
-
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     torch_dtype="auto",
@@ -51,12 +51,14 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
+# Define input query
 prompt = "Give me a short introduction to large language model."
 messages = [
     {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."},
     {"role": "user", "content": prompt}
 ]
 
+# Apply chat template
 text = tokenizer.apply_chat_template(
     messages,
     tokenize=False,
@@ -64,6 +66,7 @@ text = tokenizer.apply_chat_template(
 )
 model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
 
+# Generate rewritten query
 generated_ids = model.generate(
     **model_inputs,
     max_new_tokens=512
@@ -73,6 +76,8 @@ generated_ids = [
 ]
 
 response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+print(response)
 ```
 
 ---