1import torch
2import transformers
3from transformers import AutoTokenizer, AutoModelForCausalLM
4
5transformers.logging.set_verbosity_error()
6max_length = 512
7model_path = 'yueqingyou/BioQwen-0.5B'
8tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
9model = AutoModelForCausalLM.from_pretrained(model_path, device_map='auto', torch_dtype=torch.bfloat16, attn_implementation='flash_attention_2').eval()
10
11def predict(prompt):
12 zh_system = "你是千问生物智能助手，一个专注于生物领域的先进人工智能。"
13 en_system = "You are BioQwen, an advanced AI specializing in the field of biology."
14
15 english_count, chinese_count = 0, 0
16 for char in prompt:
17 if '\u4e00' <= char <= '\u9fff':
18 chinese_count += 1
19 elif 'a' <= char.lower() <= 'z':
20 english_count += 1
21 lang = 'zh' if chinese_count > english_count else 'en'
22
23 messages = [
24 {"role": "system", "content": zh_system if lang == 'zh' else en_system},
25 {"role": "user", "content": prompt}
26 ]
27 text = tokenizer.apply_chat_template(
28 messages,
29 tokenize=False,
30 add_generation_prompt=True
31 )
32
33 model_inputs = tokenizer([text], return_tensors="pt").to('cuda')
34
35 with torch.no_grad():
36 generated_ids = model.generate(
37 model_inputs.input_ids,
38 max_new_tokens=max_length,
39 eos_token_id=tokenizer.eos_token_id,
40 pad_token_id=tokenizer.pad_token_id,
41 do_sample=True,
42 top_p = 0.9,
43 temperature = 0.3,
44 repetition_penalty = 1.1
45 )
46
47 generated_ids = [
48 output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
49 ]
50 response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
51
52 return response.strip()
53
54prompt = '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.'
55print(f'Question:\t{prompt}\n\nAnswer:\t{predict(prompt)}')
56
For more detailed information and code, please refer to
GitHub.
