import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

class CritiqueCoreInference:
    def __init__(self, model_path):
        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).to(self.device)
        self.model.eval()

    def analyze(self, text):
        inputs = self.tokenizer(
            text, 
            return_tensors="pt", 
            padding=True, 
            truncation=True, 
            max_length=128
        ).to(self.device)

        with torch.no_grad():
            outputs = self.model(**inputs)
            probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
            conf, pred = torch.max(probs, dim=-1)

        result = "POSITIVE" if pred.item() == 1 else "NEGATIVE"
        return {
            "text": text,
            "label": result,
            "confidence": f"{conf.item() * 100:.2f}%"
        }

# Usage
if __name__ == "__main__":
    # Point this to your unzipped folder
    engine = CritiqueCoreInference("./CritiqueCore_v1_HF")
    
    sample = "The plot was a bit slow, but overall a great experience."
    prediction = engine.analyze(sample)
    print(f"Result: {prediction['label']} | Confidence: {prediction['confidence']}")