muffin2006/document-classification-env
1
1import gradio as gr2import threading3import os4import pickle5import time6import json7from flask import Flask, request, jsonify8from environment import DocumentClassificationEnv9from baseline_inference import run_task, load_or_train10from agent import TicketAgent11from rl_trainer import RLTrainer12import numpy as np13 14print("=" * 60)15print(" MetaX AI Agent - HACKATHON EDITION")16print("=" * 60)17print("Loading AI models...")18 19models = {}20for d in ["easy", "medium", "hard"]:21 models[d] = load_or_train(d)22 print(f" [OK] {d.upper()} model loaded")23 24agent = TicketAgent(model_dir=".")25print("[OK] Agent ready")26 27CATEGORIES = {28 "easy": ["General", "Billing", "Support", "Technical", "HR"],29 "medium": [30 "General",31 "Billing",32 "Support",33 "Technical",34 "HR",35 "Legal",36 "Sales",37 "Marketing",38 "Operations",39 "Complaints",40 ],41 "hard": [f"Cat_{i}" for i in range(22)],42}43 44api = Flask(__name__)45api_envs = {}46 47rl_state = {"trainer": None, "is_training": False}48 49 50@api.route("/api/reset", methods=["POST"])51def api_reset():52 data = request.json or {}53 difficulty = data.get("difficulty", "easy")54 seed = data.get("seed", 42)55 env = DocumentClassificationEnv(task_difficulty=difficulty, seed=seed)56 obs, _ = env.reset(seed=seed)57 api_envs["current"] = env58 return jsonify(59 {60 "observation": {61 k: v.tolist() if hasattr(v, "tolist") else v for k, v in obs.items()62 }63 }64 )65 66 67@api.route("/api/step", methods=["POST"])68def api_step():69 data = request.json or {}70 action = data.get("action", 0)71 env = api_envs.get("current")72 if env is None:73 return jsonify({"error": "Call /api/reset first"}), 40074 obs, reward, done, _, info = env.step(int(action))75 return jsonify(76 {77 "observation": {78 k: v.tolist() if hasattr(v, "tolist") else v for k, v in obs.items()79 },80 "reward": reward,81 "done": done,82 "info": info,83 }84 )85 86 87@api.route("/api/run", methods=["POST"])88def api_run():89 data = request.json or {}90 difficulty = data.get("difficulty", "easy")91 score, t = run_task(difficulty)92 return jsonify({"difficulty": difficulty, "score": score, "time": t})93 94 95@api.route("/api/agent/process", methods=["POST"])96def api_agent_process():97 data = request.json or {}98 text = data.get("text", "")99 difficulty = data.get("difficulty", "easy")100 ticket_id = data.get("ticket_id", None)101 if not text:102 return jsonify({"error": "text field required"}), 400103 result = agent.process_ticket(text, difficulty=difficulty, ticket_id=ticket_id)104 return jsonify(result)105 106 107@api.route("/api/agent/classify", methods=["POST"])108def api_agent_classify():109 data = request.json or {}110 text = data.get("text", "")111 difficulty = data.get("difficulty", "easy")112 category, scores = agent.classify(text, difficulty)113 priority = agent.get_priority(text)114 return jsonify({"category": category, "priority": priority, "scores": scores})115 116 117@api.route("/api/rl/train", methods=["POST"])118def api_rl_train():119 global rl_state120 data = request.json or {}121 difficulty = data.get("difficulty", "easy")122 episodes = data.get("episodes", 20)123 124 if rl_state.get("is_training"):125 return jsonify({"status": "Training already in progress"}), 400126 127 env = DocumentClassificationEnv(task_difficulty=difficulty, seed=42)128 trainer = RLTrainer(env, num_episodes=episodes)129 rl_state["trainer"] = trainer130 rl_state["is_training"] = True131 132 result = trainer.train_step()133 eval_result = trainer.evaluate(num_episodes=5)134 curve = trainer.generate_learning_curve()135 summary = trainer.get_metrics_summary()136 137 rl_state["is_training"] = False138 139 return jsonify(140 {141 "training_result": result,142 "evaluation": eval_result,143 "learning_curve": curve,144 "summary": summary,145 }146 )147 148 149@api.route("/api/rl/status", methods=["GET"])150def api_rl_status():151 if rl_state.get("trainer") is None:152 return jsonify({"status": "No training data"})153 return jsonify(rl_state["trainer"].get_metrics_summary())154 155 156def run_flask():157 api.run(host="0.0.0.0", port=7861, debug=False)158 159 160SAMPLE_TICKETS = {161 "Billing Complaint": "My invoice shows an incorrect amount! I was charged $500 but should have been charged $75. This is unacceptable!",162 "Critical Bug": "The application crashes whenever I try to upload a file larger than 10MB. This is critical!",163 "HR Policy": "I have a question about the company's maternity leave policy. How many weeks of paid leave am I entitled to?",164 "Login Issue": "I can't log into my account. The password reset link isn't working. Please help!",165 "Sales Inquiry": "I would like to inquire about upgrading our current subscription to include additional enterprise features.",166 "Legal Request": "We need legal review of the new vendor contract before signing.",167 "Frustrated Customer": "I'm extremely disappointed with the service. This is the third time this month I've had issues.",168}169 170env_state = {"env": None, "difficulty": "easy", "trainer": None}171 172 173def create_env(difficulty):174 e = DocumentClassificationEnv(task_difficulty=difficulty, seed=42)175 obs, _ = e.reset()176 env_state["env"] = e177 env_state["difficulty"] = difficulty178 model = models[difficulty]179 cats = CATEGORIES[difficulty]180 proba = model.predict_proba([obs["content"]])[0]181 top3 = sorted(enumerate(proba), key=lambda x: -x[1])[:3]182 explain = "\n".join([f" {cats[i]}: {p * 100:.1f}%" for i, p in top3])183 pred_idx = int(model.predict([obs["content"]])[0])184 pred_cat = cats[pred_idx]185 return (186 f"Environment Ready | {difficulty.upper()}",187 obs["content"],188 f"ML Prediction: {pred_cat}\n\nTop 3:\n{explain}",189 "",190 )191 192 193def classify_doc(category_idx):194 e = env_state["env"]195 if e is None:196 return "Create environment first!", "", "", ""197 obs, reward, done, _, info = e.step(int(category_idx))198 result = (199 f"{'CORRECT' if info.get('is_correct') else 'WRONG'} | Reward: {reward:.2f}"200 )201 if done:202 acc = info.get("episode_accuracy", 0)203 result += f"\nEpisode Complete! Accuracy: {acc:.1%}"204 return result, "Episode complete", "", str(info)205 difficulty = env_state["difficulty"]206 model = models[difficulty]207 cats = CATEGORIES[difficulty]208 proba = model.predict_proba([obs["content"]])[0]209 top3 = sorted(enumerate(proba), key=lambda x: -x[1])[:3]210 explain = "\n".join([f" {cats[i]}: {p * 100:.1f}%" for i, p in top3])211 pred_idx = int(model.predict([obs["content"]])[0])212 pred_cat = cats[pred_idx]213 return (214 result,215 obs["content"],216 f"ML Prediction: {pred_cat}\n\nTop 3:\n{explain}",217 str(info),218 )219 220 221def process_ticket_ui(ticket_text, difficulty):222 if not ticket_text.strip():223 return "Please enter ticket text!", "", "", "", "", ""224 result = agent.process_ticket(ticket_text, difficulty=difficulty)225 cat_out = f"{result['category']}\n Confidence: {result['confidence'] * 100:.1f}%"226 pri_out = f"{result['priority'].upper()}"227 dept_out = f"{result['department']}\n {result['email']}\n SLA: {result['sla']}"228 top3_out = "\n".join([f" {c}: {p * 100:.1f}%" for c, p in result["top3"]])229 ref_out = f"{result['ref_id']}\n {result['timestamp'][:19]}"230 reply_out = result["reply"]231 return cat_out, pri_out, dept_out, top3_out, ref_out, reply_out232 233 234def run_rl_training(difficulty, num_episodes):235 difficulty = difficulty or "easy"236 num_episodes = num_episodes or 20237 try:238 env = DocumentClassificationEnv(task_difficulty=difficulty, seed=42)239 trainer = RLTrainer(env, num_episodes=num_episodes)240 env_state["trainer"] = trainer241 result = trainer.train_step()242 curve_img = trainer.generate_learning_curve()243 summary = trainer.get_metrics_summary()244 status = f"Training Complete!\n\nResults:\n Episodes: {num_episodes}\n Avg Reward: {result.get('avg_reward', 0):.2f}\n Avg Accuracy: {result.get('avg_accuracy', 0):.1%}"245 eval_result = trainer.evaluate(num_episodes=5)246 status += f"\n\nEvaluation Accuracy: {eval_result.get('accuracy', 0):.1%}"247 return status, curve_img if curve_img else "", str(summary)248 except Exception as e:249 return f"Error: {str(e)}", "", ""250 251 252def load_sample(sample_name):253 return SAMPLE_TICKETS.get(sample_name, "")254 255 256def run_evaluation():257 rows = []258 for d in ["easy", "medium", "hard"]:259 score, t = run_task(d)260 rows.append([d.upper(), f"{score:.4f}", f"{score * 100:.1f}%", "Ready"])261 return rows262 263 264def process_image(file_obj):265 if file_obj is None:266 return "No image uploaded. Click 'Upload' button, select an image, then click 'Analyze Image'."267 return f"Image received: {type(file_obj)}. Image analysis requires PyTorch installation."268 269 270with gr.Blocks(title="MetaX AI - Hackathon Demo") as demo:271 gr.Markdown("""272 <div style="text-align: center; padding: 25px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 20px; margin-bottom: 25px;">273 <h1 style="color: white; margin: 0; font-size: 2.5em;">MetaX AI Agent</h1>274 <p style="color: white; opacity: 0.95;">Document Classification | RL Training | AI-Powered</p>275 </div>276 """)277 278 with gr.Tab("AI Agent Demo"):279 gr.Markdown("### AI-Powered Ticket Classification")280 with gr.Row():281 with gr.Column(scale=1):282 sample_dd = gr.Dropdown(283 choices=list(SAMPLE_TICKETS.keys()), label="Load Sample"284 )285 ticket_input = gr.Textbox(286 label="Ticket Text",287 lines=5,288 placeholder="Enter customer message...",289 )290 diff_agent = gr.Radio(291 ["easy", "medium", "hard"], value="medium", label="Difficulty"292 )293 btn_process = gr.Button("Process Ticket", variant="primary", size="lg")294 with gr.Column(scale=1):295 cat_out_box = gr.Textbox(label="Category", lines=2)296 pri_out_box = gr.Textbox(label="Priority", lines=2)297 dept_out_box = gr.Textbox(label="Routing", lines=2)298 top3_out_box = gr.Textbox(label="Top Predictions", lines=2)299 with gr.Row():300 ref_out_box = gr.Textbox(label="Reference", lines=1)301 reply_out_box = gr.Textbox(label="AI Response", lines=8)302 sample_dd.change(load_sample, inputs=sample_dd, outputs=ticket_input)303 btn_process.click(304 process_ticket_ui,305 inputs=[ticket_input, diff_agent],306 outputs=[307 cat_out_box,308 pri_out_box,309 dept_out_box,310 top3_out_box,311 ref_out_box,312 reply_out_box,313 ],314 )315 316 with gr.Tab("RL Training"):317 gr.Markdown("### Train RL Agent")318 with gr.Row():319 with gr.Column(scale=1):320 diff_rl = gr.Radio(321 ["easy", "medium", "hard"], value="easy", label="Difficulty"322 )323 episodes_slider = gr.Slider(10, 100, value=30, step=5, label="Episodes")324 btn_train = gr.Button("Start Training", variant="primary", size="lg")325 with gr.Column(scale=2):326 train_status = gr.Textbox(label="Status", lines=8)327 learning_curve_img = gr.Image(label="Training Progress")328 btn_train.click(329 run_rl_training,330 inputs=[diff_rl, episodes_slider],331 outputs=[train_status, learning_curve_img, train_status],332 )333 334 with gr.Tab("Multi-Modal"):335 gr.Markdown("### Image Upload (Experimental)")336 gr.Markdown(337 "Note: Full image analysis requires PyTorch. Install: `pip install torch transformers`"338 )339 with gr.Row():340 file_input = gr.File(341 label="Upload Image", file_count="single", file_types=["image"]342 )343 image_btn = gr.Button("Analyze Image")344 image_output = gr.Textbox(label="Result", lines=4)345 image_btn.click(process_image, inputs=file_input, outputs=image_output)346 347 with gr.Tab("Interactive Env"):348 gr.Markdown("### Test Classification Environment")349 with gr.Row():350 diff = gr.Radio(351 ["easy", "medium", "hard"], value="easy", label="Difficulty"352 )353 btn_create = gr.Button("Create Environment", variant="primary")354 status = gr.Textbox(label="Status")355 with gr.Row():356 doc_content = gr.Textbox(label="Document", lines=4)357 explain_box = gr.Textbox(label="Model Reasoning", lines=4)358 with gr.Row():359 category = gr.Number(label="Category Index", value=0, precision=0)360 btn_classify = gr.Button("Classify")361 result = gr.Textbox(label="Result", lines=2)362 btn_create.click(363 create_env, inputs=diff, outputs=[status, doc_content, explain_box, result]364 )365 btn_classify.click(366 classify_doc,367 inputs=category,368 outputs=[result, doc_content, explain_box, result],369 )370 371 with gr.Tab("Performance"):372 gr.Markdown("### Model Performance")373 btn_eval = gr.Button("Run Evaluation", variant="primary")374 score_table = gr.Dataframe(375 headers=["Task", "Score", "Accuracy", "Status"], label="Results"376 )377 btn_eval.click(run_evaluation, outputs=score_table)378 379 with gr.Tab("Info"):380 gr.Markdown("""381 ## API Endpoints (port 7861)382 - POST /api/reset383 - POST /api/step 384 - POST /api/agent/process385 - POST /api/rl/train386 - GET /api/rl/status387 388 ## Run389 python app_v3.py390 Then open http://localhost:7860391 """)392 393if __name__ == "__main__":394 t = threading.Thread(target=run_flask, daemon=True)395 t.start()396 demo.launch(server_name="0.0.0.0", server_port=7860)397