From e1274857765b70371148366397d32b6ebabc2711 Mon Sep 17 00:00:00 2001
From: Kannan R <kannanr@techconative.com>
Date: Thu, 2 Nov 2023 00:23:56 +0530
Subject: [PATCH] Fixed template and prompt generation

---
 generate/lora_ui_gen.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/generate/lora_ui_gen.py b/generate/lora_ui_gen.py
index 7d43bcd..d086c18 100644
--- a/generate/lora_ui_gen.py
+++ b/generate/lora_ui_gen.py
@@ -28,12 +28,12 @@
 lora_mlp = False
 lora_head = False
 
-TEMPLATE = Template('''<s>
+TEMPLATE = Template('''
 <<SYS>>The response MUST be a valid JSON. Generate UI-DSL for the below input and context.<</SYS>>
 {%- if context -%}## Input: {{ context }}{%- endif -%}
 [INS]## Instruction: {{ prompt }}[/INS]
-## Response: ```{{ response }}```
-</s>''')
+## Response: 
+''')
 
 def main(
     prompt: str = "What food do lamas eat?",
@@ -153,7 +153,7 @@ def main(
     t = time.perf_counter() - t0
 
     output = tokenizer.decode(y)
-    output = output.split("### Response:")[1].strip()
+    # output = output.split("### Response:")[1].strip()
     fabric.print(output)
 
     tokens_generated = y.size(0) - prompt_length
@@ -164,7 +164,7 @@ def main(
 def generate_prompt(example: dict) -> str:
     """Generates a standardized message to prompt the model with an instruction, optional input and a
     'response' field."""
-    return TEMPLATE.render(prompt=example['instruction'], response=example['output'], context=example.get('input', ''))
+    return TEMPLATE.render(prompt=example['instruction'], context=example.get('input', ''))
 
 if __name__ == "__main__":
     from jsonargparse import CLI