import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

class VeraMindInference:
    def __init__(self, model_path, max_len=512):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_path)
        self.model.to(self.device)
        self.model.eval()
        self.max_len = max_len

    def predict(self, text):
        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,
            max_length=self.max_len,
            return_token_type_ids=False,
            padding='max_length',
            truncation=True,
            return_attention_mask=True,
            return_tensors='pt',
        )

        input_ids = encoding['input_ids'].to(self.device)
        attention_mask = encoding['attention_mask'].to(self.device)

        with torch.no_grad():
            outputs = self.model(input_ids, attention_mask=attention_mask).logits
            prediction = torch.sigmoid(outputs).cpu().numpy()[0][0]

        is_fake = prediction >= 0.5
        confidence = prediction if is_fake else 1 - prediction

        return {
            "result": "FAKE" if is_fake else "REAL",
            "confidence": float(confidence),
            "is_fake": is_fake
        }