Update src/demo.py

src/demo.py

@@ -27,8 +27,20 @@ model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16
 # type2dataset = {}
 
 
-def generate():
-    return "1"
+def generate(input_text, sys_prompt) -> str:
+    sys_prompt = f'''[INST] <<SYS>>
+{sys_prompt}
+<</SYS>>
+
+'''
+    input_str = sys_prompt + input_text + " [/INST]"
+    
+    input_ids = tokenizer(input_str, return_tensors="pt").input_ids
+    outputs = model.generate(input_ids, max_length=512)
+
+    result = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
+
+    return result
 
 
 def random_examples(dataset_key) -> str: