Update Fluency_Score.py

Fluency_Score.py

@@ -3,13 +3,12 @@ import datasets
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 
-class Fluency_Score(evaluate.Measurement):
-    
+class FluencyScore(evaluate.Metric):
     def _info(self):
-        return evaluate.MeasurementInfo(
-            description="",
+        return evaluate.MetricInfo(
+            description="Computes the fluency score of a given text using a pre-trained model.",
             citation="",
-            inputs_description="",
+            inputs_description="A list of text strings to evaluate for fluency.",
             features=datasets.Features(
                 {
                     "texts": datasets.Value("string", id="sequence"),
@@ -17,31 +16,31 @@ class Fluency_Score(evaluate.Measurement):
             ),
             reference_urls=[],
         )
-    
-    def _download_and_prepare(self, dl_manager, device=None):
+
+    def __init__(self, device=None):
+        super().__init__()
         if device is None:
             device = "cuda" if torch.cuda.is_available() else "cpu"
-        
-        # Load the tokenizer and model from the specified repository
+        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(device)
-        self.device = device
+        self.model.to(self.device)
 
-    def _compute(self, texts, temperature=2):
-        device = self.device
-        
+    def _compute(self, texts):
+        # Tokenize the input texts
         inputs = self.tokenizer(
-            texts, 
-            return_tensors="pt", 
-            truncation=True, 
-            padding='max_length', 
+            texts,
+            return_tensors="pt",
+            truncation=True,
+            padding='max_length',
             max_length=128
-        ).to(device)
-        
-        with torch.inference_mode():
+        ).to(self.device)
+
+        # Get model predictions
+        with torch.no_grad():
             output = self.model(**inputs)
-            prediction = output.logits.clip(0, 1)
+            predictions = output.logits.clip(0, 1).squeeze().tolist()  # Convert to list
 
-        return {"classical_score": prediction}
+        return {"fluency_scores": predictions}