Update app.py

app.py

@@ -3,6 +3,7 @@ import sqlite3
 from datetime import datetime
 import gradio as gr
 from huggingface_hub import InferenceClient
+from datasets import load_dataset
 
 # ---------------------------
 # Config
@@ -10,10 +11,10 @@ from huggingface_hub import InferenceClient
 MODELS = {
     "Meta LLaMA 3.1 (8B Instruct)": "meta-llama/Llama-3.1-8B-Instruct",
     "Mistral 7B Instruct": "mistralai/Mistral-7B-Instruct-v0.3",
-    "StarCoder Backend-Focused": "bigcode/starcoderbase",
-    "CodeLlama Backend-Focused": "codellama/CodeLlama-7b-Instruct"
 }
 
+DATASETS = ["The Stack", "CodeXGLUE"]  # Dropdown for dataset selection
+
 HF_TOKEN = os.getenv("HF_TOKEN")  # Set in your Space's Secrets
 DB_PATH = "history.db"
 
@@ -23,7 +24,7 @@ SYSTEM_DEFAULT = (
     "Provide full backend code scaffolds with files, paths, and commands. "
     "Only include frontend if required for framework integration "
     "(e.g., Laravel Blade, Django templates). "
-    "If asked for unrelated frontend/UI tasks, respond politely that you are a backend assistant."
+    "If user asks for too much frontend, politely say: 'I'm a backend-focused assistant and cannot provide excessive frontend code.'"
 )
 
 # ---------------------------
@@ -96,7 +97,8 @@ def get_messages(session_id: int):
     """, (session_id,))
     rows = cur.fetchall()
     conn.close()
-    return [{"role": role, "content": content} for role, content in rows]
+    msgs = [{"role": role, "content": content} for (role, content) in rows]
+    return msgs
 
 def add_message(session_id: int, role: str, content: str):
     conn = db()
@@ -140,6 +142,13 @@ def get_client(model_choice: str):
     model_id = MODELS.get(model_choice, list(MODELS.values())[0])
     return InferenceClient(model_id, token=HF_TOKEN)
 
+def load_dataset_by_name(name: str):
+    if name == "The Stack":
+        return load_dataset("bigcode/the-stack", split="train")
+    elif name == "CodeXGLUE":
+        return load_dataset("google/code_x_glue_cc_code_to_code_trans", split="train")
+    return None
+
 # ---------------------------
 # Gradio Callbacks
 # ---------------------------
@@ -168,7 +177,7 @@ def delete_chat_cb(selected_label):
     selected = labels[0] if labels else None
     return gr.update(choices=labels, value=selected), []
 
-def send_cb(user_text, selected_label, chatbot_msgs, system_message, max_tokens, temperature, top_p, model_choice):
+def send_cb(user_text, selected_label, chatbot_msgs, system_message, max_tokens, temperature, top_p, model_choice, dataset_choice):
     sid = label_to_id(selected_label)
     if sid is None:
         sid = create_session("New chat")
@@ -178,10 +187,21 @@ def send_cb(user_text, selected_label, chatbot_msgs, system_message, max_tokens,
     add_message(sid, "user", user_text)
     update_session_title_if_needed(sid, user_text)
 
+    # Load selected dataset (used internally for reference / semantic checks)
+    dataset = load_dataset_by_name(dataset_choice)
+
     api_messages = build_api_messages(sid, system_message)
     display_msgs = get_messages(sid)
     display_msgs.append({"role": "assistant", "content": ""})
 
+    # Block excessive frontend requests
+    if "frontend" in user_text.lower() and "too much" in user_text.lower():
+        warning = "⚠️ I'm a backend-focused assistant and cannot provide excessive frontend code."
+        display_msgs[-1]["content"] = warning
+        yield (display_msgs, "", selected_label)
+        add_message(sid, "assistant", warning)
+        return
+
     client = get_client(model_choice)
     partial = ""
     try:
@@ -203,7 +223,7 @@ def send_cb(user_text, selected_label, chatbot_msgs, system_message, max_tokens,
         display_msgs[-1]["content"] = err
         yield (display_msgs, "", selected_label)
 
-def regenerate_cb(selected_label, system_message, max_tokens, temperature, top_p, model_choice):
+def regenerate_cb(selected_label, system_message, max_tokens, temperature, top_p, model_choice, dataset_choice):
     sid = label_to_id(selected_label)
     if sid is None:
         return [], ""
@@ -225,6 +245,7 @@ def regenerate_cb(selected_label, system_message, max_tokens, temperature, top_p
         conn.close()
         msgs = get_messages(sid)
 
+    dataset = load_dataset_by_name(dataset_choice)
     api_messages = [{"role": "system", "content": system_message.strip()}] + msgs
     display_msgs = msgs + [{"role": "assistant", "content": ""}]
 
@@ -258,25 +279,16 @@ if not labels:
     labels, _ = list_sessions()
 default_selected = labels[0] if labels else None
 
-with gr.Blocks(title="Backend-Focused CRUD Assistant", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="Backend-Focused LLaMA/Mistral CRUD Assistant", theme=gr.themes.Soft()) as demo:
     gr.HTML("""
     <style>
-      button {
-        background-color: #22c55e !important;
-        color: #ffffff !important;
-        border: none !important;
-      }
-      button:hover {
-        background-color: #16a34a !important;
-      }
-      button:focus {
-        outline: 2px solid #166534 !important;
-        outline-offset: 2px;
-      }
+      button { background-color: #22c55e !important; color: #ffffff !important; border: none !important; }
+      button:hover { background-color: #16a34a !important; }
+      button:focus { outline: 2px solid #166534 !important; outline-offset: 2px; }
     </style>
     """)
 
-    gr.Markdown("## 🗄️ Backend-Focused CRUD Automation — with Persistent History")
+    gr.Markdown("## 🗄️ LLaMA & Mistral Backend-Focused CRUD Automation — with Persistent History")
 
     with gr.Row():
         with gr.Column(scale=1, min_width=260):
@@ -300,11 +312,19 @@ with gr.Blocks(title="Backend-Focused CRUD Assistant", theme=gr.themes.Soft()) a
                 interactive=True
             )
 
+            gr.Markdown("### 📚 Dataset Selection")
+            dataset_choice = gr.Dropdown(
+                choices=DATASETS,
+                value=DATASETS[0],
+                label="Select a dataset",
+                interactive=True
+            )
+
             gr.Markdown("### ⚙️ Generation Settings")
             system_box = gr.Textbox(
                 value=SYSTEM_DEFAULT,
                 label="System message",
-                lines=4
+                lines=5
             )
             max_tokens = gr.Slider(256, 4096, value=1200, step=16, label="Max tokens")
             temperature = gr.Slider(0.0, 2.0, value=0.25, step=0.05, label="Temperature")
@@ -318,7 +338,6 @@ with gr.Blocks(title="Backend-Focused CRUD Assistant", theme=gr.themes.Soft()) a
                 send_btn = gr.Button("Send ▶️", variant="primary")
                 regen_btn = gr.Button("Regenerate 🔁", variant="secondary")
 
-    # --- Button interactions ---
     refresh_btn.click(refresh_sessions_cb, outputs=session_list)
     new_btn.click(new_chat_cb, outputs=[session_list, chatbot, user_box])
     del_btn.click(delete_chat_cb, inputs=session_list, outputs=[session_list, chatbot])
@@ -326,19 +345,19 @@ with gr.Blocks(title="Backend-Focused CRUD Assistant", theme=gr.themes.Soft()) a
 
     send_btn.click(
         send_cb,
-        inputs=[user_box, session_list, chatbot, system_box, max_tokens, temperature, top_p, model_choice],
+        inputs=[user_box, session_list, chatbot, system_box, max_tokens, temperature, top_p, model_choice, dataset_choice],
         outputs=[chatbot, user_box, session_list]
     )
 
     user_box.submit(
         send_cb,
-        inputs=[user_box, session_list, chatbot, system_box, max_tokens, temperature, top_p, model_choice],
+        inputs=[user_box, session_list, chatbot, system_box, max_tokens, temperature, top_p, model_choice, dataset_choice],
         outputs=[chatbot, user_box, session_list]
     )
 
     regen_btn.click(
         regenerate_cb,
-        inputs=[session_list, system_box, max_tokens, temperature, top_p, model_choice],
+        inputs=[session_list, system_box, max_tokens, temperature, top_p, model_choice, dataset_choice],
         outputs=chatbot
     )