Update README.md

README.md

@@ -18,8 +18,11 @@ For model inference, please refer to the following example code:
 
 ```python
 import torch
+import transformers
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
+transformers.logging.set_verbosity_error()
+max_length = 512
 model_path = 'yueqingyou/BioQwen-1.8B'
 tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
 model = AutoModelForCausalLM.from_pretrained(model_path, device_map='auto', torch_dtype=torch.bfloat16, attn_implementation='flash_attention_2').eval()
@@ -57,7 +60,7 @@ def predict(prompt):
             do_sample=True,
             top_p = 0.9,
             temperature = 0.3,
-            repetition_penalty = 1.0
+            repetition_penalty = 1.1
         )
 
     generated_ids = [
@@ -68,7 +71,8 @@ def predict(prompt):
     return response.strip()
 
 prompt = 'I am suffering from irregular periods. I am currently taking medication Levothyroxine 50. My T3 is 0.87 ng/mL, T4 is 8.30 ug/dL, TSH is 2.43 uIU/mL. I am 34 years old, weigh 75 kg, and 5 feet tall. Please advice.'
-predict(prompt)
+print(f'Question:\t{prompt}\n\nAnswer:\t{predict(prompt)}')
+
 ```
 
 For more detailed information and code, please refer to [GitHub](https://github.com/yueqingyou/BioQwen).