from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

class BertInference:
    def __init__(self, model_path):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model = AutoModelForSequenceClassification.from_pretrained(model_path).to(self.device)
        self.tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
        # self.tokenizer = AutoTokenizer.from_pretrained("FacebookAI/xlm-roberta-base")
        self.label_map = {
            0: "query_with_pdf",
            1: "summarize_pdf",
            2: "query_metadata"
        }

    def predict(self, text):
        # Tokenize
        inputs = self.tokenizer(
            text,
            return_tensors="pt"
        ).to(self.device)

        # Get prediction
        with torch.no_grad():
            outputs = self.model(**inputs)
            predictions = torch.softmax(outputs.logits, dim=1)
            predicted_class = torch.argmax(predictions, dim=1).item()
            confidence = predictions[0][predicted_class].item()

        return {
            "predicted_class": self.label_map[predicted_class],
            "confidence": confidence,
            "all_probabilities": {
                self.label_map[i]: prob.item()
                for i, prob in enumerate(predictions[0])
            }
        }

def main():
    # Initialize the model
    model_path = "output_dir_decision"  # Path to your saved model
    # model_path = "output_xlm_roberta_bert"
    inferencer = BertInference(model_path)

    # Example usage
    test_questions = [
        "What is television",
        "What is the summary",
        "What is GPU",
        "What is the title of this pdf?"
    ]

    for question in test_questions:
        result = inferencer.predict(question)
        print(f"\nQuestion: {question}")
        print(f"Predicted Class: {result['predicted_class']}")
        print(f"Confidence: {result['confidence']:.4f}")
        print("All Probabilities:")
        for class_name, prob in result['all_probabilities'].items():
            print(f"  {class_name}: {prob:.4f}")

if __name__ == "__main__":
    main()