Update README.md

README.md

@@ -342,37 +342,20 @@ import torch
 model_name = "baichuan-inc/Baichuan-M1-14B-Base"  
 tokenizer = AutoTokenizer.from_pretrained(model_name,trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(model_name,trust_remote_code=True,torch_dtype = torch.bfloat16).cuda()
-# 2. Input prompt text
-prompt = "May I ask you some questions about medical knowledge?"
-
-# 3. Encode the input text for the model
-messages = [
-    {"role": "system", "content": "You are a helpful assistant."},
-    {"role": "user", "content": prompt}
-]
-text = tokenizer.apply_chat_template(
-    messages,
-    tokenize=False,
-    add_generation_prompt=True
-)
-model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
-
-# 4. Generate text
-generated_ids = model.generate(
-    **model_inputs,
-    max_new_tokens=512
-)
-generated_ids = [
-    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
-]
-
-# 5. Decode the generated text
-response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-
-
-# 6. Output the result
-print("Generated text:")
-print(response)
+
+input_text = "I have recently recovered from my cold."
+
+
+inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+
+outputs = model.generate(
+        inputs["input_ids"],
+        max_length=100, 
+    )
+
+generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+print("Generated Text:")
+print(generated_text)
 ```
 
 ---