BnSa3d/coaching-api

1import torch2import json3import re4from flask import Flask, request, jsonify5from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig6from peft import PeftModel7 8# =========================9# CONFIGURATION10# =========================11BASE_MODEL   = "Qwen/Qwen2.5-3B-Instruct"12ADAPTER_PATH = "./Adapter_Coaching"13 14# Determine the best available hardware (MPS for Mac, CUDA for Cloud, CPU as fallback)15if torch.cuda.is_available():16    device = "cuda"17elif torch.backends.mps.is_available():18    device = "mps"19else:20    device = "cpu"21 22# =========================23# MODEL LOADING (Optimized for RAM)24# =========================25print(f"Detected device: {device}")26print("Loading tokenizer...")27tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)28 29# 4-bit quantization configuration to ensure the model runs on consumer hardware/Mac30bnb_config = BitsAndBytesConfig(31    load_in_4bit=True,32    bnb_4bit_use_double_quant=True,33    bnb_4bit_quant_type="nf4",34    bnb_4bit_compute_dtype=torch.float1635)36 37print("Loading base model...")38base_model = AutoModelForCausalLM.from_pretrained(39    BASE_MODEL,40    quantization_config=bnb_config if device != "cuda" else None,  # Use 4-bit on Mac/Local, Full on Cloud GPU41    torch_dtype=torch.float16,42    device_map="auto" if device == "cuda" else {"": "cpu"},43    trust_remote_code=True44)45 46print("Loading LoRA adapter...")47model = PeftModel.from_pretrained(base_model, ADAPTER_PATH)48 49# Move model to target device if not already handled by device_map50if device == "mps":51    model = model.to("mps")52 53model.eval()54print("Agent Coaching Model and Adapter loaded successfully!")55 56# =========================57# COACHING PROMPT58# =========================59SYSTEM_MESSAGE = """\60You are an expert AI Customer Support Coach.61 62Your task is to analyze a customer–agent conversation and provide structured coaching feedback in Egyptian Arabic.63 64STRICT RULES:65- Output MUST be valid JSON66- Output MUST follow the exact schema below67- Do NOT include explanations, comments, or markdown68- Do NOT add extra keys69- Use Egyptian Arabic for all text values70- Output JSON only71 72SCHEMA:73{74  "ai_recommendations": string,75  "weakness_analysis": string,76  "suggested_learning": string,77  "encouragement_quote": string78}79 80GUIDELINES:81- ai_recommendations: bullet-pointed coaching tips (use * per bullet)82- weakness_analysis: specific weaknesses observed in the agent's behaviour83- suggested_learning: concrete skills or techniques the agent should study84- encouragement_quote: a single motivational line to uplift the agent85 86Return ONLY the JSON object.87"""88 89# =========================90# UTILITIES91# =========================92 93def build_conversation_text(conversation):94    """95    Accepts either:96      - A list of dicts: [{"agent": "..."}, {"customer": "..."}]97      - A plain string (already formatted)98    Returns a formatted AGENT / CUSTOMER text block.99    """100    if isinstance(conversation, list):101        lines = []102        for msg in conversation:103            role, text = list(msg.items())[0]104            lines.append(f"{role.upper()} : {text}")105        return "\n".join(lines)106    return conversation  # already a plain string107 108 109def build_messages(conversation_text):110    return [111        {"role": "system", "content": SYSTEM_MESSAGE},112        {"role": "user",   "content": conversation_text},113    ]114 115 116def apply_chat_format(messages):117    return tokenizer.apply_chat_template(118        messages,119        tokenize=False,120        add_generation_prompt=True121    )122 123 124def clean_output(text):125    """126    Extracts and parses JSON from the model's raw string response.127    """128    text = text.strip()129    try:130        match = re.search(r"\{.*\}", text, re.DOTALL)131        if match:132            return json.loads(match.group())133    except Exception as e:134        return {"error": "JSON Parse Error", "details": str(e), "raw_output": text}135    return {"raw_output": text}136 137 138def analyze_agent(conversation_text):139    """140    Processes the conversation through the coaching model141    and returns structured Arabic coaching feedback.142    """143    messages = build_messages(conversation_text)144    prompt   = apply_chat_format(messages)145 146    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)147 148    with torch.no_grad():149        outputs = model.generate(150            **inputs,151            max_new_tokens=1024,152            temperature=0.1,153            top_p=0.9,154            do_sample=False,155            pad_token_id=tokenizer.eos_token_id156        )157 158    # Extract only the newly generated tokens (Assistant response)159    input_length     = inputs.input_ids.shape[1]160    generated_tokens = outputs[0][input_length:]161    decoded          = tokenizer.decode(generated_tokens, skip_special_tokens=True)162 163    return clean_output(decoded)164 165# =========================166# FLASK API ENDPOINTS167# =========================168app = Flask(__name__)169 170 171@app.route("/health", methods=["GET"])172def health():173    return jsonify({"status": "active", "device": device, "model": "Agent Coaching"})174 175 176@app.route("/coach", methods=["POST"])177def coach():178    """179    Expects JSON body with a 'conversation' field.180 181    Accepted formats:182      1) List of dicts:183         {"conversation": [{"agent": "Hello"}, {"customer": "Hi"}]}184 185      2) Plain formatted string:186         {"conversation": "AGENT : Hello\nCUSTOMER : Hi"}187    """188    data = request.get_json()189 190    if not data or "conversation" not in data:191        return jsonify({"error": "Missing 'conversation' field"}), 400192 193    conversation = data["conversation"]194 195    # Validate type196    if not isinstance(conversation, (list, str)):197        return jsonify({"error": "'conversation' must be a list of message dicts or a plain string"}), 400198 199    # Build readable text from whatever format was sent200    try:201        conversation_text = build_conversation_text(conversation)202    except Exception as e:203        return jsonify({"error": "Failed to parse conversation", "details": str(e)}), 400204 205    result = analyze_agent(conversation_text)206    return jsonify(result)207 208 209# =========================210# ENTRY POINT211# =========================212if __name__ == "__main__":213    # Port 8080 is standard for many cloud deployments214    app.run(host="0.0.0.0", port=7860)215