Create inference.py

inference.py

@@ -0,0 +1,62 @@
+import torch
+from transformers import (
+    RobertaForSequenceClassification,
+    DebertaForSequenceClassification,
+    RobertaTokenizer,
+    DebertaTokenizer,
+    RobertaConfig,
+    DebertaConfig
+)
+
+class EnsembleInference:
+    def __init__(self, model_path, device='cpu'):
+        self.device = device
+        self.roberta_tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
+        self.deberta_tokenizer = DebertaTokenizer.from_pretrained("microsoft/deberta-base")
+        self.load_models(model_path)
+
+    def load_models(self, path):
+        state = torch.load(path, map_location=self.device)
+
+        roberta_config = RobertaConfig.from_dict(state['model_configs']['roberta_config'])
+        deberta_config = DebertaConfig.from_dict(state['model_configs']['deberta_config'])
+
+        self.roberta_model = RobertaForSequenceClassification(roberta_config).to(self.device)
+        self.deberta_model = DebertaForSequenceClassification(deberta_config).to(self.device)
+
+        self.roberta_model.load_state_dict(state['roberta_state_dict'])
+        self.deberta_model.load_state_dict(state['deberta_state_dict'])
+
+        self.roberta_model.eval()
+        self.deberta_model.eval()
+
+    def predict(self, text):
+        roberta_inputs = self.roberta_tokenizer(text, return_tensors='pt', truncation=True, padding=True, max_length=512)
+        deberta_inputs = self.deberta_tokenizer(text, return_tensors='pt', truncation=True, padding=True, max_length=512)
+
+        roberta_inputs = {k: v.to(self.device) for k, v in roberta_inputs.items()}
+        deberta_inputs = {k: v.to(self.device) for k, v in deberta_inputs.items()}
+
+        with torch.no_grad():
+            roberta_logits = self.roberta_model(**roberta_inputs).logits.squeeze()
+            deberta_logits = self.deberta_model(**deberta_inputs).logits.squeeze()
+
+            roberta_prob = torch.sigmoid(roberta_logits).item()
+            deberta_prob = torch.sigmoid(deberta_logits).item()
+
+            avg_prob = (roberta_prob + deberta_prob) / 2
+            is_ai = avg_prob > 0.5
+            prediction = "AI generated" if is_ai else "Human written"
+
+            roberta_conf = roberta_prob if is_ai else 1 - roberta_prob
+            deberta_conf = deberta_prob if is_ai else 1 - deberta_prob
+            avg_conf = avg_prob if is_ai else 1 - avg_prob
+
+        return {
+            'prediction': prediction,
+            'confidence': f"{avg_conf:.2%}",
+            'details': {
+                'roberta_confidence': f"{roberta_conf:.2%}",
+                'deberta_confidence': f"{deberta_conf:.2%}"
+            }
+        }