Update README.md

README.md

@@ -34,12 +34,14 @@ config = PeftConfig.from_pretrained("DeathReaper0965/Qwen2.5-3B-Inst-SQL-Reasoni
 
 model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-3B-Instruct", max_length=2560)
 model = PeftModel.from_pretrained(model, "DeathReaper0965/Qwen2.5-3B-Inst-SQL-Reasoning-GRPO", is_trainable=False)
+
 tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B-Instruct", max_length = 2560)
 
-prompt = [
-  {
-    'role': 'system',
-    'content': """\
+def create_prompt(schemas, question):
+    prompt = [
+      {
+        'role': 'system',
+        'content': """\
 You are an expert SQL Query Writer. 
 Given relevant Schemas and the Question, you first understand the problem entirely and then reason about the best possible approach to come up with an answer.
 Once, you are confident in your reasoning, you will then start generating the SQL Query as the answer that accurately solves the given question leveraging some or all schemas.
@@ -58,13 +60,26 @@ FROM TABLE_NAME
 WHERE
     CONDITION
 </answer>"""
-  },
-  {
-    'role': 'user',
-    'content': """\
+      },
+      {
+        'role': 'user',
+        'content': f"""\
 SCHEMAS:
 ---------------
 
+{schemas}
+
+---------------
+
+QUESTION: "{question}"\
+"""
+      }
+    ]
+
+    return prompt
+
+
+schemas = """\
 CREATE TABLE lab (
     subject_id text,
     hadm_id text,
@@ -127,23 +142,23 @@ CREATE TABLE prescriptions (
     formulary_drug_cd text,
     route text,
     drug_dose text
-)
+)\
+"""
 
----------------
+question = "How many patients whose admission type is emergency and diagnoses icd9 code is 56210?"
 
-QUESTION: "how many patients whose admission type is emergency and diagnoses icd9 code is 56210?"
-"""
-  }
-]
+example_prompt = create_prompt(schemas, question)
+
+streamer = TextStreamer(tokenizer, skip_prompt=True)
 
-inputs = tokenizer.apply_chat_template(prompt,
+inputs = tokenizer.apply_chat_template(example_prompt,
                                        tokenize=True,
                                        add_generation_prompt=True,
                                        return_dict=True,
                                        return_tensors="pt")
 
 with torch.inference_mode():
-    outputs = model.generate(**inputs, max_new_tokens=256)
+    outputs = model.generate(**inputs, max_new_tokens=1024, streamer=streamer)
 
 outputs = tokenizer.batch_decode(outputs)
 print(outputs[0].split("<|im_start|>assistant")[-1])