import evaluate
import datasets
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

class FluencyScore(evaluate.Metric):
    def _info(self):
        return evaluate.MetricInfo(
            description="Computes the fluency score of a given text using a pre-trained model.",
            citation="",
            inputs_description="A list of text strings to evaluate for fluency.",
            features=datasets.Features(
                {
                    "texts": datasets.Value("string", id="sequence"),
                }
            ),
            reference_urls=[],
        )

    def __init__(self, device=None):
        super().__init__()
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        self.device = device

        # Load the tokenizer and model
        self.tokenizer = AutoTokenizer.from_pretrained("Baleegh/Fluency_Score")
        self.model = AutoModelForSequenceClassification.from_pretrained("Baleegh/Fluency_Score")
        self.model.to(self.device)

    def _compute(self, texts):
        # Tokenize the input texts
        inputs = self.tokenizer(
            texts,
            return_tensors="pt",
            truncation=True,
            padding='max_length',
            max_length=128
        ).to(self.device)

        # Get model predictions
        with torch.no_grad():
            output = self.model(**inputs)
            predictions = output.logits.clip(0, 1).squeeze().tolist()  # Convert to list

        return {"fluency_scores": predictions}