Upload app.py

app.py

@@ -32,7 +32,11 @@ def generate_lean4_code(prompt):
     #response, history = model.chat(tokenizer, prompt, history=[], meta_instruction="")
     #print(response, history)
     #return response
-    input_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt").to(device)
+    chat = [
+        {"role": "system", "content": "You are a Lean4 expert who can write good Lean4 code based on natural language mathematical theorem and proof"},
+        {"role": "user", "content": prompt},
+    ]
+    input_ids = tokenizer.apply_chat_template(chat, return_tensors="pt").to(device)
     results = model.generate(input_ids, max_new_tokens=1024, eos_token_id=terminators, do_sample=True, temperature=0.85, top_p=0.9)
     result_str = tokenizer.decode(results[0], skip_special_tokens=True)
     return result_str