File size: 11,038 Bytes
24f95f0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
"""
Domain Classifier Trainer for Janus self-improvement system.

Trains a domain classifier using curated examples from the observation layer.
The classifier predicts which domain a query belongs to (finance, technology, etc.)
"""

import json
import logging
import os
from pathlib import Path
from typing import Dict, List, Tuple, Optional
from collections import defaultdict
import math

logger = logging.getLogger(__name__)

# Domain list from query_classifier.py
DOMAINS = [
    "finance",
    "technology",
    "healthcare",
    "policy",
    "science",
    "geopolitics",
    "energy",
    "critical_thinking",
    "emotional_intelligence",
    "philosophy",
    "business",
    "education",
    "general",
]


class DomainClassifierTrainer:
    """
    Trains a domain classifier using curated examples.
    Uses Naive Bayes approach for simplicity and interpretability.
    """

    def __init__(self, data_dir: Optional[Path] = None):
        if data_dir is None:
            try:
                from app.config import DATA_DIR as BASE_DATA_DIR
            except ImportError:
                BASE_DATA_DIR = Path(__file__).resolve().parent.parent.parent / "data"

            self.data_dir = Path(BASE_DATA_DIR) / "curation"
        else:
            self.data_dir = data_dir

        self.curated_file = self.data_dir / "curated_examples.jsonl"

        # Model parameters
        self.domain_word_counts = defaultdict(lambda: defaultdict(int))
        self.domain_total_words = defaultdict(int)
        self.domain_doc_counts = defaultdict(int)
        self.vocab = set()
        self.total_docs = 0

        # Load existing model if available
        self._load_model()

    def _load_model(self):
        """Load pre-trained model parameters if they exist."""
        model_file = self.data_dir / "domain_classifier_model.json"
        if model_file.exists():
            try:
                with open(model_file) as f:
                    model_data = json.load(f)
                    self.domain_word_counts = defaultdict(
                        lambda: defaultdict(int),
                        model_data.get("domain_word_counts", {}),
                    )
                    self.domain_total_words = defaultdict(
                        int, model_data.get("domain_total_words", {})
                    )
                    self.domain_doc_counts = defaultdict(
                        int, model_data.get("domain_doc_counts", {})
                    )
                    self.vocab = set(model_data.get("vocab", []))
                    self.total_docs = model_data.get("total_docs", 0)
                logger.info(f"Loaded domain classifier model from {model_file}")
            except Exception as e:
                logger.error(f"Failed to load model: {e}")

    def _save_model(self):
        """Save model parameters to disk."""
        model_file = self.data_dir / "domain_classifier_model.json"
        try:
            model_data = {
                "domain_word_counts": dict(self.domain_word_counts),
                "domain_total_words": dict(self.domain_total_words),
                "domain_doc_counts": dict(self.domain_doc_counts),
                "vocab": list(self.vocab),
                "total_docs": self.total_docs,
            }
            with open(model_file, "w") as f:
                json.dump(model_data, f, indent=2)
            logger.info(f"Saved domain classifier model to {model_file}")
        except Exception as e:
            logger.error(f"Failed to save model: {e}")

    def _tokenize(self, text: str) -> List[str]:
        """Simple tokenization - lowercase and split on non-alphanumeric."""
        import re

        # Convert to lowercase and split on non-alphanumeric characters
        tokens = re.findall(r"\b\w+\b", text.lower())
        return tokens

    def train_from_curated_examples(self) -> Dict[str, any]:
        """
        Train the domain classifier using curated examples.
        Returns training statistics.
        """
        if not self.curated_file.exists():
            logger.warning(f"No curated examples found at {self.curated_file}")
            return {"error": "No training data available"}

        # Reset counters
        self.domain_word_counts = defaultdict(lambda: defaultdict(int))
        self.domain_total_words = defaultdict(int)
        self.domain_doc_counts = defaultdict(int)
        self.vocab = set()
        self.total_docs = 0

        # Process each curated example
        try:
            with open(self.curated_file) as f:
                for line_num, line in enumerate(f, 1):
                    if line.strip():
                        try:
                            example = json.loads(line)
                            query = example.get("query", "")
                            domain = example.get("domain", "general")

                            # Skip if domain not in our list
                            if domain not in DOMAINS:
                                domain = "general"

                            # Tokenize the query
                            tokens = self._tokenize(query)

                            # Update counts
                            self.domain_doc_counts[domain] += 1
                            self.total_docs += 1

                            for token in tokens:
                                self.domain_word_counts[domain][token] += 1
                                self.domain_total_words[domain] += 1
                                self.vocab.add(token)

                        except json.JSONDecodeError as e:
                            logger.error(f"Invalid JSON on line {line_num}: {e}")
                            continue

            # Calculate and save model
            self._save_model()

            stats = {
                "total_documents": self.total_docs,
                "vocabulary_size": len(self.vocab),
                "domain_distribution": dict(self.domain_doc_counts),
                "training_complete": True,
            }

            logger.info(f"Domain classifier training complete: {stats}")
            return stats

        except Exception as e:
            logger.error(f"Training failed: {e}")
            return {"error": str(e)}

    def predict_domain(self, query: str) -> Tuple[str, float]:
        """
        Predict the domain for a given query.
        Returns (domain, confidence) tuple.
        """
        if self.total_docs == 0:
            return "general", 0.0

        tokens = self._tokenize(query)
        if not tokens:
            return "general", 0.0

            # Calculate log probabilities for each domain
        log_probs = {}
        vocab_size = len(self.vocab)

        for domain in DOMAINS:
            # Prior probability P(domain) - handle unseen domains
            if self.domain_doc_counts[domain] == 0:
                # If domain not seen in training, use a small probability
                prior = math.log(1e-10)
            else:
                prior = math.log(self.domain_doc_counts[domain] / self.total_docs)

            # Likelihood P(query|domain) = product of P(word|domain) for each word
            likelihood = 0.0
            for token in tokens:
                # Laplace smoothing: P(word|domain) = (count(word in domain) + 1) / (total_words_in_domain + vocab_size)
                word_count = self.domain_word_counts[domain].get(token, 0)
                total_words_in_domain = self.domain_total_words[domain]
                prob = (word_count + 1) / (total_words_in_domain + vocab_size)
                likelihood += math.log(prob)

            log_probs[domain] = prior + likelihood

        # Convert log probabilities to probabilities
        max_log_prob = max(log_probs.values())
        probs = {
            domain: math.exp(log_prob - max_log_prob)
            for domain, log_prob in log_probs.items()
        }

        # Normalize to get probabilities
        prob_sum = sum(probs.values())
        if prob_sum > 0:
            probs = {domain: prob / prob_sum for domain, prob in probs.items()}
        else:
            # Uniform distribution if something went wrong
            probs = {domain: 1.0 / len(DOMAINS) for domain in DOMAINS}

        # Get the domain with highest probability
        predicted_domain = max(probs, key=probs.get)
        confidence = probs[predicted_domain]

        return predicted_domain, confidence

    def evaluate(self) -> Dict[str, any]:
        """
        Evaluate the classifier on the curated examples.
        Returns accuracy and other metrics.
        """
        if not self.curated_file.exists() or self.total_docs == 0:
            return {"error": "No training data or model not trained"}

        correct = 0
        total = 0
        domain_stats = defaultdict(lambda: {"correct": 0, "total": 0})

        try:
            with open(self.curated_file) as f:
                for line in f:
                    if line.strip():
                        example = json.loads(line)
                        query = example.get("query", "")
                        actual_domain = example.get("domain", "general")

                        if actual_domain not in DOMAINS:
                            actual_domain = "general"

                        predicted_domain, confidence = self.predict_domain(query)

                        if predicted_domain == actual_domain:
                            correct += 1
                            domain_stats[actual_domain]["correct"] += 1

                        domain_stats[actual_domain]["total"] += 1
                        total += 1

            accuracy = correct / total if total > 0 else 0.0

            # Calculate per-domain accuracy
            domain_accuracies = {}
            for domain, stats in domain_stats.items():
                if stats["total"] > 0:
                    domain_accuracies[domain] = stats["correct"] / stats["total"]
                else:
                    domain_accuracies[domain] = 0.0

            return {
                "accuracy": accuracy,
                "correct_predictions": correct,
                "total_predictions": total,
                "domain_accuracies": domain_accuracies,
                "domain_distribution": dict(self.domain_doc_counts),
            }

        except Exception as e:
            logger.error(f"Evaluation failed: {e}")
            return {"error": str(e)}


# Global instance
domain_classifier_trainer = DomainClassifierTrainer()


def train_domain_classifier() -> Dict[str, any]:
    """Convenience function to train the domain classifier."""
    return domain_classifier_trainer.train_from_curated_examples()


def predict_domain(query: str) -> Tuple[str, float]:
    """Convenience function to predict domain for a query."""
    return domain_classifier_trainer.predict_domain(query)


def evaluate_domain_classifier() -> Dict[str, any]:
    """Convenience function to evaluate the domain classifier."""
    return domain_classifier_trainer.evaluate()