TanujInsane/document-classification-env
3
1import gradio as gr2import threading3import os4import pickle5from flask import Flask, request, jsonify6from environment import DocumentClassificationEnv7from baseline_inference import run_task, load_or_train8from agent import TicketAgent9 10print("Pre-training models...")11models = {}12for d in ["easy", "medium", "hard"]:13 models[d] = load_or_train(d)14 print(f" {d} model ready!")15print("All models ready!")16 17agent = TicketAgent(model_dir=".")18print("Agent ready!")19 20CATEGORIES = {21 "easy": ["General","Billing","Support","Technical","HR"],22 "medium": ["General","Billing","Support","Technical","HR","Legal","Sales","Marketing","Operations","Complaints"],23 "hard": [f"Cat_{i}" for i in range(22)]24}25 26api = Flask(__name__)27api_envs = {}28 29@api.route("/api/reset", methods=["POST"])30def api_reset():31 data = request.json or {}32 difficulty = data.get("difficulty", "easy")33 seed = data.get("seed", 42)34 env = DocumentClassificationEnv(task_difficulty=difficulty, seed=seed)35 obs, _ = env.reset(seed=seed)36 api_envs["current"] = env37 return jsonify({"observation": {k: v.tolist() if hasattr(v, 'tolist') else v for k, v in obs.items()}})38 39@api.route("/api/step", methods=["POST"])40def api_step():41 data = request.json or {}42 action = data.get("action", 0)43 env = api_envs.get("current")44 if env is None:45 return jsonify({"error": "Call /api/reset first"}), 40046 obs, reward, done, _, info = env.step(int(action))47 return jsonify({48 "observation": {k: v.tolist() if hasattr(v, 'tolist') else v for k, v in obs.items()},49 "reward": reward, "done": done, "info": info50 })51 52@api.route("/api/run", methods=["POST"])53def api_run():54 data = request.json or {}55 difficulty = data.get("difficulty", "easy")56 score, t = run_task(difficulty)57 return jsonify({"difficulty": difficulty, "score": score, "time": t})58 59@api.route("/api/agent/process", methods=["POST"])60def api_agent_process():61 data = request.json or {}62 text = data.get("text", "")63 difficulty = data.get("difficulty", "easy")64 ticket_id = data.get("ticket_id", None)65 if not text:66 return jsonify({"error": "text field required"}), 40067 result = agent.process_ticket(text, difficulty=difficulty, ticket_id=ticket_id)68 return jsonify(result)69 70@api.route("/api/agent/classify", methods=["POST"])71def api_agent_classify():72 data = request.json or {}73 text = data.get("text", "")74 difficulty = data.get("difficulty", "easy")75 category, scores = agent.classify(text, difficulty)76 priority = agent.get_priority(text)77 return jsonify({"category": category, "priority": priority, "scores": scores})78 79def run_flask():80 api.run(host="0.0.0.0", port=7861, debug=False)81 82env_state = {"env": None, "difficulty": "easy"}83 84def create_env(difficulty):85 e = DocumentClassificationEnv(task_difficulty=difficulty, seed=42)86 obs, _ = e.reset()87 env_state["env"] = e88 env_state["difficulty"] = difficulty89 model = models[difficulty]90 cats = CATEGORIES[difficulty]91 proba = model.predict_proba([obs["content"]])[0]92 top3 = sorted(enumerate(proba), key=lambda x: -x[1])[:3]93 explain = "\n".join([f" {cats[i] if i < len(cats) else f'Cat_{i}'}: {p*100:.1f}%" for i,p in top3])94 pred_idx = int(model.predict([obs["content"]])[0])95 pred_cat = cats[pred_idx] if pred_idx < len(cats) else f"Cat_{pred_idx}"96 return (f"✅ Environment created! Difficulty: {difficulty}",97 obs["content"],98 f"🤖 Model predicts: **{pred_cat}**\n\nTop 3 predictions:\n{explain}",99 "")100 101def classify_doc(category_idx):102 e = env_state["env"]103 if e is None:104 return "Create environment first!", "", "", ""105 obs, reward, done, _, info = e.step(int(category_idx))106 result = f"Reward: {reward:.2f} | Correct: {info.get('is_correct', False)}"107 if done:108 acc = info.get('episode_accuracy', 0)109 result += f"\n🏁 Episode done! Final Accuracy: {acc:.2%}"110 return result, "Episode complete — create new environment", "", str(info)111 difficulty = env_state["difficulty"]112 model = models[difficulty]113 cats = CATEGORIES[difficulty]114 proba = model.predict_proba([obs["content"]])[0]115 top3 = sorted(enumerate(proba), key=lambda x: -x[1])[:3]116 explain = "\n".join([f" {cats[i] if i < len(cats) else f'Cat_{i}'}: {p*100:.1f}%" for i,p in top3])117 pred_idx = int(model.predict([obs["content"]])[0])118 pred_cat = cats[pred_idx] if pred_idx < len(cats) else f"Cat_{pred_idx}"119 return (result, obs["content"],120 f"🤖 Model predicts: **{pred_cat}**\n\nTop 3 predictions:\n{explain}",121 str(info))122 123def run_baseline_all():124 rows = []125 for d in ["easy", "medium", "hard"]:126 score, t = run_task(d)127 rows.append([d.upper(), f"{score:.4f}", f"{score*100:.1f}%", f"{t:.1f}s"])128 return rows129 130def process_ticket_ui(ticket_text, difficulty):131 if not ticket_text.strip():132 return "❌ Please enter ticket text!", "", "", "", "", ""133 result = agent.process_ticket(ticket_text, difficulty=difficulty)134 category_out = f"🏷️ {result['category']} (confidence: {result['confidence']*100:.1f}%)"135 priority_out = f"🚨 {result['priority'].upper()}"136 dept_out = f"🏢 {result['department']}\n📧 {result['email']}\n⏱️ SLA: {result['sla']}"137 top3_out = "\n".join([f" {c}: {p*100:.1f}%" for c,p in result['top3']])138 ref_out = f"🎫 {result['ref_id']} | ⏰ {result['timestamp'][:19]}"139 reply_out = result['reply']140 return category_out, priority_out, dept_out, top3_out, ref_out, reply_out141 142SAMPLE_TICKETS = {143 "Billing complaint": "My invoice shows an incorrect amount. I was charged $150 but should have been charged $75. Please review and correct this immediately.",144 "Bug report": "The application crashes whenever I try to upload a file larger than 10MB. This is a critical issue affecting my workflow.",145 "HR query": "I have a question about the company's maternity leave policy. How many weeks of paid leave am I entitled to?",146 "Chat message": "Hey, I can't log into my account. The password reset link isn't working either. Please help ASAP!",147 "Email ticket": "Dear Support, I would like to inquire about upgrading my current subscription plan to include additional users.",148}149 150def load_sample(sample_name):151 return SAMPLE_TICKETS.get(sample_name, "")152 153with gr.Blocks(title="Document Classification OpenEnv") as demo:154 gr.Markdown("# 📄 Document Classification OpenEnv")155 gr.Markdown("Real-world customer support ticket routing environment for RL agent training.")156 157 with gr.Tab("🤖 Agent Demo"):158 gr.Markdown("### Full Automation — ML Classification + Rule-based Priority + Auto-Reply")159 with gr.Row():160 with gr.Column(scale=2):161 sample_dd = gr.Dropdown(162 choices=list(SAMPLE_TICKETS.keys()),163 label="📋 Load Sample Ticket",164 value=None165 )166 ticket_input = gr.Textbox(167 label="✍️ Ticket / Message Text",168 lines=6,169 placeholder="Paste email, chat message, bug report, HR query..."170 )171 diff_agent = gr.Radio(["easy","medium","hard"], value="medium", label="Model Difficulty")172 btn_process = gr.Button("🚀 Process Ticket", variant="primary")173 with gr.Column(scale=2):174 ref_out_box = gr.Textbox(label="🎫 Reference ID & Timestamp")175 cat_out_box = gr.Textbox(label="🏷️ Category & Confidence")176 pri_out_box = gr.Textbox(label="🚨 Priority")177 dept_out_box = gr.Textbox(label="🏢 Routed To", lines=3)178 top3_out_box = gr.Textbox(label="📊 Top 3 Predictions", lines=3)179 reply_out_box = gr.Textbox(label="📧 Auto-Generated Reply", lines=12)180 181 sample_dd.change(load_sample, inputs=sample_dd, outputs=ticket_input)182 btn_process.click(183 process_ticket_ui,184 inputs=[ticket_input, diff_agent],185 outputs=[cat_out_box, pri_out_box, dept_out_box, top3_out_box, ref_out_box, reply_out_box]186 )187 188 with gr.Tab("🎮 Interactive Demo"):189 diff = gr.Radio(["easy","medium","hard"], value="easy", label="Difficulty")190 btn_create = gr.Button("Create Environment", variant="primary")191 status = gr.Textbox(label="Status")192 with gr.Row():193 doc_content = gr.Textbox(label="📝 Document Content", lines=6)194 explain_box = gr.Textbox(label="🧠 Explainability — Model Reasoning", lines=6)195 category = gr.Number(label="Category Index (easy: 0-4, medium: 0-9, hard: 0-21)", value=0)196 btn_classify = gr.Button("Classify Document")197 result = gr.Textbox(label="Result")198 info_box = gr.Textbox(label="Info")199 btn_create.click(create_env, inputs=diff, outputs=[status, doc_content, explain_box, result])200 btn_classify.click(classify_doc, inputs=category, outputs=[result, doc_content, explain_box, info_box])201 202 with gr.Tab("📊 Baseline Evaluation"):203 gr.Markdown("TF-IDF + Logistic Regression baseline — models pre-trained at startup for instant results.")204 btn_eval = gr.Button("Run All Tasks", variant="primary")205 score_table = gr.Dataframe(headers=["Task","Score","Accuracy","Time"], label="Results")206 btn_eval.click(run_baseline_all, outputs=score_table)207 208 with gr.Tab("📋 Environment Info"):209 gr.Markdown("""210## Environment Design211**Task**: Classify customer support documents into correct departments.212 213| Difficulty | Categories | Episodes | Reward |214|------------|-----------|----------|--------|215| Easy | 5 | 20 | +1.0 correct, -0.4 wrong |216| Medium | 10 | 30 | +1.0 correct, -0.4 wrong |217| Hard | 22 | 50 | +1.0 correct, -0.4 wrong |218 219## Agent Pipeline220`Ticket Input` → `ML Classify` → `Rule Priority` → `Department Route` → `Auto Reply`221 222## API Endpoints (port 7861)223- `POST /api/reset` — `{"difficulty": "easy", "seed": 42}`224- `POST /api/step` — `{"action": 0}`225- `POST /api/agent/process` — `{"text": "my invoice is wrong", "difficulty": "medium"}`226- `POST /api/agent/classify` — `{"text": "urgent bug", "difficulty": "easy"}`227 """)228 229if __name__ == "__main__":230 t = threading.Thread(target=run_flask, daemon=True)231 t.start()232 demo.launch(server_name="0.0.0.0", server_port=7860)233