Views
No views yet
1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
3model = AutoModelForCausalLM.from_pretrained("jinee/note", load_in_4bit=True, device_map="auto")
4tokenizer = AutoTokenizer.from_pretrained("jinee/note")
5tokenizer.padding_side = 'right'
6tokenizer.add_eos_token = True
7tokenizer.pad_token = tokenizer.eos_token
8tokenizer.add_eos_token, tokenizer.add_bos_token
9
10instruction = '''
11As a doctor, you need to create a discharge summary based on input data.
12Never change the dates or numbers in the input data and use them as is. And please follow the format below for your report.
13Also, never make up information that is not in the input data, and write a report only with information that can be identified from the input data.
14
151. Patient information (SUBJECT_ID, HADM_ID, hospitalization and discharge date, hospitalization period, gender, date of birth, age, allergy)
162. Diagnostic information and past history (if applicable)
173. Surgery or procedure information
184. Significant medication administration during hospitalization and discharge medication history
195. Meaningful lab tests during hospitalization
206. Summary of significant text records/notes
217. Discharge outcomes and treatment plan
228. Overall summary of at least 500 characters in lines including the above contents
23'''
24torch.cuda.empty_cache()
25
26def generation(model, tokenizer, input_data):
27 pipe = pipeline('text-generation',
28 model = model,
29 tokenizer = tokenizer,
30 torch_dtype=torch.bfloat16,
31 device_map = 'auto')
32 global instruction
33
34 sequences = pipe(
35 f"[INST]{instruction}: {input_data} [/INST]",
36 do_sample=True,
37 max_new_tokens=1024,
38 temperature=0.7,
39 top_k=50,
40 top_p=0.95,
41 early_stopping =True,
42 num_return_sequences=1,)
43
44 text = sequences[0]['generated_text']
45 start_index = text.find('[/INST]')
46 if start_index != -1:
47 summary_ = text[start_index + len('[/INST]'):]
48 return(summary_)
49 else:
50 return("'[summary_] 'is not founded.")
51| Parameter | SFT | DPO |
|---|---|---|
| r | 16 | 16 |
| lora alpha | 16 | 16 |
| lora dropout | 0.05 | 0.05 |
| target | q, k, v, o, gate | q, k, v, o, gate |
| Parameter | SFT | DPO |
|---|---|---|
| early stopping patience | 3 | 3 |
| early stopping threshold | 0.0005 | 0.0005 |
| train epochs | 20 | 3 |
| per device train batch size | 4 | 1 |
| per device eval batch size | 8 (default) | 1 |
| optimizer | paged adamw 8bit | paged adamw 8bit |
| lr scheduler | cosine | cosine |
| wramup ratio | 0.3 | 0.1 |
| gradient accumulation step | 2 | 2 |
| evaluation strategy | step | step |
| eval step | 10 | 5 |