Update app.py

app.py

@@ -1,26 +1,13 @@
-# import gradio as gr
-
-# def greet(name):
-#     return "Hello " + name
-
-# demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-# demo.launch()
-
-
-"""
 import gradio as gr
+import json, time, os
+from huggingface_hub import HfApi
 from llama_cpp import Llama
 
-# Modello .gguf ospitato nel repo del modello
-MODEL_REPO = "fedealex/iris"
-MODEL_FILE = "model-1b-q8_0.gguf"
-
-
-
-import os
-
-HF_TOKEN = os.getenv('HF_TOKEN')
-hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "iris-flags")
+MODEL_REPO = "fedealex/llama-1B"       # The model is on hf
+MODEL_FILE = "model-1b-q8_0.gguf"      # Name of the model file
+DATASET_REPO = "fedealex/flags"        # Flages saved on hf
+HF_TOKEN = os.getenv("HF_TOKEN")       # To access hf
+LOCAL_FLAGS = "flags.json"             # You must save locally before push to hf
 
 
 
@@ -35,18 +22,10 @@ llm = Llama.from_pretrained(
     temperature=0.7
 )
 
-#def chat(message, history):
-    # costruisci contesto della conversazione
-    #prompt = ""
-    #for human, bot in history:
-    #    prompt += f"<|user|>{human}\n<|assistant|>{bot}\n"
-    #prompt += f"<|user|>{message}\n<|assistant|>"
-
-    #output = llm(prompt, max_tokens=350)
-    #return output["choices"][0]["text"].strip()
+### The Chat Model
+def chat_model(message, history):
 
-def chat(message, history):
-    # construct conversation context
+    # Retrieve the context
     prompt = ""
     for item in history:
         role = item["role"]
@@ -54,86 +33,64 @@ def chat(message, history):
         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()
 
-demo = gr.ChatInterface(
-    fn=chat,
-    title="Iris – Fine-Tuned LLM",
-    flagging_mode="manual",
-    flagging_dir="evals",
-    flagging_options=["Like", "Dislike", "Incorrect", "Inappropriate", "Gibberish", "Spam", "Other"],
-    flagging_callback=hf_writer,
-    save_history=True
-)
-
-demo.launch()
-"""
 
-import gradio as gr
-import json, time, os
-from huggingface_hub import HfApi
-
-DATASET_ID = "fedealex/flags"        # your dataset
-SECRET_TOKEN = os.getenv("HF_TOKEN")   # token stored in your Space secrets
-LOCAL_JSONL = "flags.jsonl"   # local buffer file before upload
-
-
-# -------------------------------------------------
-# Example Chat Model
-# -------------------------------------------------
-def chat_model(message, history):
-    return message[::-1]  # dummy: reverse text
 
-
-# -------------------------------------------------
-# Save flags to Hugging Face Dataset
-# -------------------------------------------------
+### Save the flags
 def save_flag_to_dataset(history, reason):
-    # Prepare flag record
+    # The record to be submitted
     record = {
         "timestamp": time.time(),
         "history": history,
         "reason": reason
     }
 
-    # Append locally to JSONL
-    with open(LOCAL_JSONL, "a") as f:
+    # First we save it locally
+    with open(LOCAL_FLAGS, "a") as f:
         f.write(json.dumps(record) + "\n")
 
-    # Upload to the dataset (overwrite the JSONL with the new one)
+    # Then we send to the hf dataset
     api = HfApi()
     api.upload_file(
-        path_or_fileobj=LOCAL_JSONL,
-        path_in_repo=LOCAL_JSONL,   # same name in the dataset
-        repo_id=DATASET_ID,
+        path_or_fileobj=LOCAL_FLAGS,
+        path_in_repo=LOCAL_FLAGS,
+        repo_id=DATASET_REPO,
         repo_type="dataset",
-        token=SECRET_TOKEN
+        token=HF_TOKEN
     )
 
-    return "Flag successfully saved to dataset!"
+    if reason=="GOOD":
+        return "Feedback reported successfully!"
+    else:
+        return "Flag reported successfully!"
+
 
 
-# -------------------------------------------------
-# UI
-# -------------------------------------------------
+### The Gradio App
 with gr.Blocks() as app:
-    gr.Markdown("### Chatbot with Multi-Reason Flagging → HF Dataset")
+    # Title
+    gr.Markdown("### Scalable Machine Learning Lab 2")
 
-    # Chat UI
-    chat = gr.ChatInterface(
+    # Chat Box
+    chat_box = gr.ChatInterface(
         fn=chat_model,
         chatbot=gr.Chatbot(height=400),
-        textbox=gr.Textbox(placeholder="Say something..."),
-        title="Chatbot",
-        description="Flags saved to the dn/flags dataset."
+        textbox=gr.Textbox(placeholder="How can I help you today?"),
+        title="Llama Finetuned",
+        description="You are using the model: "+MODEL_REPO+"/"+MODEL_FILE
     )
 
-    # Flagging UI components
+    # Feedback Buttons
+    gr.Markdown("### Let us know what do you think of our chatbot!")
+    good_btn = gr.Button("Appreciate conversation ❤", variant="huggingface")
     flag_btn = gr.Button("Flag Conversation", variant="stop")
 
-    with gr.Group(visible=False) as flag_modal:
-        gr.Markdown("### Why are you flagging this conversation?")
+    # We allow the user to select flagging reason
+    with gr.Group(visible=False) as flag_group:
+        gr.Markdown("### What kind of problem are you facing?")
         reason_dd = gr.Dropdown(
             choices=[
                 "Offensive / Toxic",
@@ -143,27 +100,41 @@ with gr.Blocks() as app:
                 "Biased Output",
                 "Other"
             ],
-            label="Reason"
+            label="Flagging Reason"
         )
-        submit_flag = gr.Button("Submit Flag", variant="primary")
-        cancel_flag = gr.Button("Cancel")
+        submit_flag_btn = gr.Button("Submit Flag", variant="primary")
+        cancel_flag_btn = gr.Button("Cancel")
 
-    flag_status = gr.Textbox(label="Flag Status", visible=True)
+    # To inform the user about feedback status
+    feedback_status = gr.Textbox(label="Feedback Status", visible=True)
 
-    # Show modal
-    flag_btn.click(lambda: gr.update(visible=True), None, flag_modal)
+    # Button callbacks
+    flag_btn.click(
+        lambda: gr.update(visible=True),
+        inputs=None,
+        outputs=flag_group
+    )
 
-    # Hide modal
-    cancel_flag.click(lambda: gr.update(visible=False), None, flag_modal)
+    cancel_flag_btn.click(
+        lambda: gr.update(visible=False),
+        inputs=None,
+        outputs=flag_group
+    )
 
-    # Submit → save to dataset → hide modal
-    submit_flag.click(
+    submit_flag_btn.click(
         lambda history, reason: save_flag_to_dataset(history, reason),
-        inputs=[chat.chatbot, reason_dd],
-        outputs=flag_status
+        inputs=[chat_box.chatbot, reason_dd],
+        outputs=feedback_status
     ).then(
-        lambda: gr.update(visible=False), None, flag_modal
+        lambda: gr.update(visible=False), None, flag_group
     )
+    
+    good_btn.click(
+        lambda history, reason: save_flag_to_dataset(history, reason),
+        inputs=[chat_box.chatbot, "GOOD"],
+        outputs=feedback_status
+    )
+
 
 
-app.launch()
+app.launch()