update 2 endpoints

app.py

@@ -1,30 +1,85 @@
+import json
+import os
 import gradio as gr
 from typing import List, Dict, Any
 
-from ml_service import get_ml_service
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
 
 
-def predict_labels(text: str) -> List[Dict[str, Any]]:
-    service = get_ml_service()
-    return service.predict(text)
+class DualModelService:
+    def __init__(self, intent_dir: str = "./intent_model", tone_dir: str = "./tone_model"):
+        self.intent_dir = intent_dir
+        self.tone_dir = tone_dir
+        self._clf_cache: Dict[str, Any] = {}
+        self._labels_cache: Dict[str, List[str]] = {}
 
+    def _load_labels(self, model_dir: str) -> List[str]:
+        if model_dir in self._labels_cache:
+            return self._labels_cache[model_dir]
+        label_path = os.path.join(model_dir, "label_cols.json")
+        with open(label_path, "r", encoding="utf-8") as f:
+            labels = json.load(f)
+        self._labels_cache[model_dir] = labels
+        return labels
 
-title = "ML Text Classification (DistilBERT)"
-description = (
-    "Enter text to get multi-label predictions. Uses the model in `model/`."
-)
+    def _load_pipeline(self, model_dir: str):
+        if model_dir in self._clf_cache:
+            return self._clf_cache[model_dir]
+        model = AutoModelForSequenceClassification.from_pretrained(model_dir)
+        tokenizer = AutoTokenizer.from_pretrained(model_dir)
+        clf = pipeline(
+            "text-classification",
+            model=model,
+            tokenizer=tokenizer,
+            return_all_scores=True,
+        )
+        self._clf_cache[model_dir] = clf
+        return clf
+
+    def _predict_sorted(self, text: str, model_dir: str) -> List[Dict[str, float]]:
+        clf = self._load_pipeline(model_dir)
+        labels = self._load_labels(model_dir)
+        processed = text.replace("|||", "[SEP]")
+        scores = clf(processed)[0]
+        pairs = [(labels[idx], scores[idx]["score"]) for idx in range(len(labels))]
+        pairs_sorted = sorted(pairs, key=lambda x: x[1], reverse=True)
+        return [{"label": name, "score": float(score)} for name, score in pairs_sorted]
+
+    def predict_intent(self, text: str) -> List[Dict[str, float]]:
+        return self._predict_sorted(text, self.intent_dir)
+
+    def predict_tone(self, text: str) -> List[Dict[str, float]]:
+        return self._predict_sorted(text, self.tone_dir)
+
+
+_dual_service = DualModelService()
+
+
+def predict_intent(text: str) -> List[Dict[str, float]]:
+    return _dual_service.predict_intent(text)
+
+
+def predict_tone(text: str) -> List[Dict[str, float]]:
+    return _dual_service.predict_tone(text)
+
+
+title = "Text Classification: Intent & Tone"
+description = "Enter text; get sorted labels with scores for Intent and Tone."
 
 with gr.Blocks(title=title) as demo:
     gr.Markdown(f"# {title}\n{description}")
-    with gr.Row():
-        inp = gr.Textbox(lines=6, label="Input Text", placeholder="Type text here…")
-    with gr.Row():
-        out = gr.JSON(label="Predictions")
-    with gr.Row():
-        btn = gr.Button("Predict")
-
-    # Expose predict API at /predict
-    btn.click(predict_labels, inputs=inp, outputs=out, api_name="predict")
+    with gr.Tabs():
+        with gr.TabItem("Intent"):
+            intent_in = gr.Textbox(lines=6, label="Input Text", placeholder="Type text here…")
+            intent_out = gr.JSON(label="Intent Predictions (sorted)")
+            intent_btn = gr.Button("Predict Intent")
+            intent_btn.click(predict_intent, inputs=intent_in, outputs=intent_out, api_name="/intent")
+
+        with gr.TabItem("Tone"):
+            tone_in = gr.Textbox(lines=6, label="Input Text", placeholder="Type text here…")
+            tone_out = gr.JSON(label="Tone Predictions (sorted)")
+            tone_btn = gr.Button("Predict Tone")
+            tone_btn.click(predict_tone, inputs=tone_in, outputs=tone_out, api_name="/tone")
 
 
 if __name__ == "__main__":

intent_model/config.json

@@ -0,0 +1,42 @@
+{
+  "activation": "gelu",
+  "architectures": [
+    "DistilBertForSequenceClassification"
+  ],
+  "attention_dropout": 0.1,
+  "dim": 768,
+  "dropout": 0.1,
+  "dtype": "float32",
+  "hidden_dim": 3072,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1",
+    "2": "LABEL_2",
+    "3": "LABEL_3",
+    "4": "LABEL_4",
+    "5": "LABEL_5",
+    "6": "LABEL_6"
+  },
+  "initializer_range": 0.02,
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1,
+    "LABEL_2": 2,
+    "LABEL_3": 3,
+    "LABEL_4": 4,
+    "LABEL_5": 5,
+    "LABEL_6": 6
+  },
+  "max_position_embeddings": 512,
+  "model_type": "distilbert",
+  "n_heads": 12,
+  "n_layers": 6,
+  "pad_token_id": 0,
+  "problem_type": "multi_label_classification",
+  "qa_dropout": 0.1,
+  "seq_classif_dropout": 0.2,
+  "sinusoidal_pos_embds": false,
+  "tie_weights_": true,
+  "transformers_version": "4.57.1",
+  "vocab_size": 30522
+}

intent_model/label_cols.json

@@ -0,0 +1,9 @@
+[
+  "partner_intent_distance",
+  "partner_intent_flirt",
+  "partner_intent_invite",
+  "partner_intent_rapport",
+  "partner_intent_reframe",
+  "partner_intent_smalltalk",
+  "partner_intent_test"
+]

intent_model/special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
+}

intent_model/tokenizer_config.json

@@ -0,0 +1,56 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "101": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "102": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "103": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "[CLS]",
+  "do_lower_case": true,
+  "extra_special_tokens": {},
+  "mask_token": "[MASK]",
+  "model_max_length": 512,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "DistilBertTokenizer",
+  "unk_token": "[UNK]"
+}

tone_model/config.json

@@ -0,0 +1,46 @@
+{
+  "activation": "gelu",
+  "architectures": [
+    "DistilBertForSequenceClassification"
+  ],
+  "attention_dropout": 0.1,
+  "dim": 768,
+  "dropout": 0.1,
+  "dtype": "float32",
+  "hidden_dim": 3072,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1",
+    "2": "LABEL_2",
+    "3": "LABEL_3",
+    "4": "LABEL_4",
+    "5": "LABEL_5",
+    "6": "LABEL_6",
+    "7": "LABEL_7",
+    "8": "LABEL_8"
+  },
+  "initializer_range": 0.02,
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1,
+    "LABEL_2": 2,
+    "LABEL_3": 3,
+    "LABEL_4": 4,
+    "LABEL_5": 5,
+    "LABEL_6": 6,
+    "LABEL_7": 7,
+    "LABEL_8": 8
+  },
+  "max_position_embeddings": 512,
+  "model_type": "distilbert",
+  "n_heads": 12,
+  "n_layers": 6,
+  "pad_token_id": 0,
+  "problem_type": "multi_label_classification",
+  "qa_dropout": 0.1,
+  "seq_classif_dropout": 0.2,
+  "sinusoidal_pos_embds": false,
+  "tie_weights_": true,
+  "transformers_version": "4.57.1",
+  "vocab_size": 30522
+}

tone_model/label_cols.json

@@ -0,0 +1,11 @@
+[
+  "partner_tone_assertive",
+  "partner_tone_curious",
+  "partner_tone_plain",
+  "partner_tone_playful",
+  "partner_tone_provocative",
+  "partner_tone_suggestive",
+  "partner_tone_uncertain",
+  "partner_tone_vulnerable",
+  "partner_tone_warm"
+]

tone_model/special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
+}

tone_model/tokenizer_config.json

@@ -0,0 +1,56 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "101": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "102": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "103": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "[CLS]",
+  "do_lower_case": true,
+  "extra_special_tokens": {},
+  "mask_token": "[MASK]",
+  "model_max_length": 512,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "DistilBertTokenizer",
+  "unk_token": "[UNK]"
+}