|
|
|
|
|
import torch |
|
|
from transformers import AutoTokenizer, AutoModelForSequenceClassification |
|
|
import torch.nn as nn |
|
|
from typing import List, Dict, Any |
|
|
import numpy as np |
|
|
|
|
|
class EnhancedProgressiveAMDModel(nn.Module): |
|
|
"""Enhanced model that incorporates utterance count information""" |
|
|
|
|
|
def __init__(self, base_model_name: str, utterance_embedding_dim: int = 8): |
|
|
super().__init__() |
|
|
|
|
|
|
|
|
self.bert = AutoModelForSequenceClassification.from_pretrained( |
|
|
base_model_name, num_labels=1 |
|
|
) |
|
|
|
|
|
|
|
|
self.utterance_count_embedding = nn.Embedding(4, utterance_embedding_dim) |
|
|
|
|
|
|
|
|
bert_hidden_size = self.bert.config.hidden_size |
|
|
self.enhanced_classifier = nn.Sequential( |
|
|
nn.Linear(bert_hidden_size + utterance_embedding_dim, 64), |
|
|
nn.ReLU(), |
|
|
nn.Dropout(0.1), |
|
|
nn.Linear(64, 1) |
|
|
) |
|
|
|
|
|
self.bert.classifier = nn.Identity() |
|
|
|
|
|
def forward(self, input_ids, attention_mask, utterance_count=None): |
|
|
bert_outputs = self.bert.bert(input_ids=input_ids, attention_mask=attention_mask) |
|
|
pooled_output = bert_outputs.pooler_output |
|
|
|
|
|
if utterance_count is not None: |
|
|
utterance_emb = self.utterance_count_embedding(utterance_count) |
|
|
combined_features = torch.cat([pooled_output, utterance_emb], dim=1) |
|
|
logits = self.enhanced_classifier(combined_features) |
|
|
else: |
|
|
batch_size = pooled_output.size(0) |
|
|
zero_utterance_emb = torch.zeros(batch_size, 8, device=pooled_output.device) |
|
|
combined_features = torch.cat([pooled_output, zero_utterance_emb], dim=1) |
|
|
logits = self.enhanced_classifier(combined_features) |
|
|
|
|
|
return logits |
|
|
|
|
|
class ProductionEnhancedAMDClassifier: |
|
|
"""Production-ready enhanced AMD classifier""" |
|
|
|
|
|
def __init__(self, model_path: str, tokenizer_name: str = 'prajjwal1/bert-tiny', device: str = 'auto'): |
|
|
if device == 'auto': |
|
|
if torch.backends.mps.is_available(): |
|
|
self.device = torch.device('mps') |
|
|
elif torch.cuda.is_available(): |
|
|
self.device = torch.device('cuda') |
|
|
else: |
|
|
self.device = torch.device('cpu') |
|
|
else: |
|
|
self.device = torch.device(device) |
|
|
|
|
|
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) |
|
|
self.model = EnhancedProgressiveAMDModel(tokenizer_name) |
|
|
self.model.load_state_dict(torch.load(model_path, map_location=self.device)) |
|
|
self.model.to(self.device) |
|
|
self.model.eval() |
|
|
|
|
|
self.max_length = 128 |
|
|
self.threshold = 0.5 |
|
|
|
|
|
print(f"Enhanced AMD classifier loaded on {self.device}") |
|
|
|
|
|
def extract_user_utterances(self, transcript: List[Dict[str, Any]]) -> List[str]: |
|
|
user_utterances = [] |
|
|
for utterance in transcript: |
|
|
if utterance.get("speaker", "").lower() == "user": |
|
|
content = utterance.get("content", "").strip() |
|
|
if content: |
|
|
user_utterances.append(content) |
|
|
return user_utterances |
|
|
|
|
|
@torch.no_grad() |
|
|
def predict(self, transcript: List[Dict[str, Any]]) -> Dict[str, Any]: |
|
|
user_utterances = self.extract_user_utterances(transcript) |
|
|
|
|
|
if not user_utterances: |
|
|
return { |
|
|
'prediction': 'Human', |
|
|
'machine_probability': 0.0, |
|
|
'confidence': 0.5, |
|
|
'utterance_count': 0 |
|
|
} |
|
|
|
|
|
utt1 = user_utterances[0] if len(user_utterances) >= 1 else "" |
|
|
utt2 = user_utterances[1] if len(user_utterances) >= 2 else "" |
|
|
utt3 = user_utterances[2] if len(user_utterances) >= 3 else "" |
|
|
|
|
|
combined_text = " ".join([utt for utt in [utt1, utt2, utt3] if utt.strip()]) |
|
|
utterance_count = min(len(user_utterances), 3) |
|
|
|
|
|
encoding = self.tokenizer( |
|
|
combined_text, |
|
|
add_special_tokens=True, |
|
|
max_length=self.max_length, |
|
|
padding='max_length', |
|
|
truncation=True, |
|
|
return_tensors='pt' |
|
|
) |
|
|
|
|
|
input_ids = encoding['input_ids'].to(self.device) |
|
|
attention_mask = encoding['attention_mask'].to(self.device) |
|
|
utterance_count_tensor = torch.tensor([utterance_count], dtype=torch.long).to(self.device) |
|
|
|
|
|
logits = self.model( |
|
|
input_ids=input_ids, |
|
|
attention_mask=attention_mask, |
|
|
utterance_count=utterance_count_tensor |
|
|
) |
|
|
|
|
|
machine_prob = torch.sigmoid(logits.squeeze(-1)).item() |
|
|
prediction = 'Machine' if machine_prob >= self.threshold else 'Human' |
|
|
confidence = max(machine_prob, 1 - machine_prob) |
|
|
|
|
|
return { |
|
|
'prediction': prediction, |
|
|
'machine_probability': machine_prob, |
|
|
'confidence': confidence, |
|
|
'utterance_count': utterance_count, |
|
|
'available_utterances': len(user_utterances) |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
