Views
No views yet
1import onnxruntime_genai as og
2
3model = og.Model('soap5_onnx')
4tokenizer = og.Tokenizer(model)
5tokenizer_stream = tokenizer.create_stream()
6
7# Search options - exact match to original
8search_options = {
9 'max_length': 4096,
10 'temperature': 0.1,
11 'top_p': 0.9,
12 'do_sample': True,
13 'batch_size': 1
14}
15
16soap_note_prompt = """You are an expert medical professor assisting in the creation of medically accurate SOAP summaries.
17Please ensure the response follows the structured format: S:, O:, A:, P: without using markdown or special formatting.
18Create a Medical SOAP note summary from the dialogue, following these guidelines:\n
19S (Subjective): Summarize the patient's reported symptoms, including chief complaint and relevant history.
20Rely on the patient's statements as the primary source and ensure standardized terminology.\n
21O (Objective): Highlight critical findings such as vital signs, lab results, and imaging, emphasizing important details like the side of the body affected and specific dosages.
22Include normal ranges where relevant.\n
23A (Assessment): Offer a concise assessment combining subjective and objective data. State the primary diagnosis and any differential diagnoses, noting potential complications and the prognostic outlook.\n
24P (Plan): Outline the management plan, covering medication, diet, consultations, and education. Ensure to mention necessary referrals to other specialties and address compliance challenges.\n
25Considerations: Compile the report based solely on the transcript provided. Use concise medical jargon and abbreviations for effective doctor communication.\n
26Please format the summary in a clean, simple list format without using markdown or bullet points. Use 'S:', 'O:', 'A:', 'P:' directly followed by the text. Avoid any styling or special characters.
27TRANSCRIPT: \n"""
28
29text = input("Input: ")
30if not text:
31 print("Error, input cannot be empty")
32 exit()
33
34# Method 1: Force generation by adding a SOAP starter after the prompt
35full_prompt = soap_note_prompt + text
36
37# Use the most complete Llama format
38chat_template = "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n{prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\nS: "
39
40prompt = chat_template.format(prompt=full_prompt)
41
42input_tokens = tokenizer.encode(prompt)
43print(f"Tokens in prompt: {len(input_tokens)}")
44
45params = og.GeneratorParams(model)
46params.set_search_options(**search_options)
47generator = og.Generator(model, params)
48generator.append_tokens(input_tokens)
49
50print("\nGenerating SOAP note...")
51print("S: ", end='', flush=True) # We already have "S: " in the prompt
52
53# Generate the rest of the SOAP note
54generated_text = ""
55token_count = 0
56
57try:
58 while not generator.is_done() and token_count < 2000: # Limit to 2000 tokens for safety
59 generator.generate_next_token()
60 new_token = generator.get_next_tokens()[0]
61 decoded = tokenizer_stream.decode(new_token)
62
63 # Skip if we're still in the input echo phase
64 if token_count < 50 and (text[:20] in generated_text + decoded):
65 token_count += 1
66 continue
67
68 print(decoded, end='', flush=True)
69 generated_text += decoded
70 token_count += 1
71
72 # Stop if we see end markers
73 if any(marker in decoded for marker in ["<|eot_id|>", "<|end_of_text|>", "</s>"]):
74 break
75
76except KeyboardInterrupt:
77 print("\nInterrupted")
78
79print()
80
81# If that didn't work, try Method 2: Different prompt structure
82if len(generated_text.strip()) < 50 or text[:50] in generated_text:
83 print("\n\nMethod 1 didn't work well. Trying alternative method...")
84
85 del generator # Clean up
86
87 # Try a simpler approach - maybe the model expects a different format
88 simple_prompt = f"{soap_note_prompt}{text}\n\nSOAP Note:\nS: "
89
90 input_tokens = tokenizer.encode(simple_prompt)
91
92 params = og.GeneratorParams(model)
93 params.set_search_options(**search_options)
94 generator = og.Generator(model, params)
95 generator.append_tokens(input_tokens)
96
97 print("\nGenerating with simplified format...")
98 print("S: ", end='', flush=True)
99
100 generated_text = ""
101 token_count = 0
102
103 try:
104 while not generator.is_done() and token_count < 2000:
105 generator.generate_next_token()
106 new_token = generator.get_next_tokens()[0]
107 decoded = tokenizer_stream.decode(new_token)
108
109 print(decoded, end='', flush=True)
110 generated_text += decoded
111 token_count += 1
112
113 if any(marker in decoded for marker in ["<|eot_id|>", "<|end_of_text|>", "</s>"]):
114 break
115
116 except KeyboardInterrupt:
117 print("\nInterrupted")
118
119 print()
120 del generator
121
122print("\n--- Generation Complete ---")
123'''