Update app.py

app.py

@@ -13,9 +13,29 @@ LOCAL_FLAGS = "flags.json"             # You must save locally before push to hf
 
 print("Loading model...")
 
+llm = Llama.from_pretrained(
+    repo_id=MODEL_REPO,
+    filename=MODEL_FILE,
+    n_threads=2,
+    n_batch=128,
+    n_ctx=2048,
+    temperature=0.7
+)
+
 ### The Chat Model
 def chat_model(message, history):
-    return message[::-1]  # dummy: reverse text
+
+    # Retrieve the context
+    prompt = ""
+    for item in history:
+        role = item["role"]
+        text = item["content"][0]["text"]
+        prompt += f"<|{role}|>{text}\n"
+    prompt += f"<|user|>{message}\n<|assistant|>"
+
+    # Invoke the model
+    output = llm(prompt, max_tokens=350)
+    return output["choices"][0]["text"].strip()
 
 
 
@@ -111,7 +131,7 @@ with gr.Blocks() as app:
     
     good_btn.click(
         lambda history, reason: save_flag_to_dataset(history, reason),
-        inputs=[chat_box.chatbot, "GOOD"],
+        inputs=[chat_box.chatbot, gr.Markdown("GOOD")],
         outputs=feedback_status
     )