update

app_regression.py

@@ -14,8 +14,8 @@ from models.conversation import conv_templates
 from typing import List
 
 
-processor = AutoProcessor.from_pretrained("Mantis-VL/mantis-8b-idefics2-video-eval-20k-mantis-2epoch_4096_regression")
-model = Idefics2ForSequenceClassification.from_pretrained("Mantis-VL/mantis-8b-idefics2-video-eval-20k-mantis-2epoch_4096_regression", torch_dtype=torch.bfloat16).eval()
+processor = AutoProcessor.from_pretrained("Mantis-VL/mantis-8b-idefics2-video-eval-40k-mantis-2epoch_4096_regression")
+model = Idefics2ForSequenceClassification.from_pretrained("Mantis-VL/mantis-8b-idefics2-video-eval-40k-mantis-2epoch_4096_regression", torch_dtype=torch.bfloat16).eval()
 
 MAX_NUM_FRAMES = 24
 conv_template = conv_templates["idefics_2"]
@@ -64,14 +64,13 @@ all the frames of video are as follows:
 """
 
 
-aspect_mapping={
-    1:"visual quality",
-    2:"object consistency",
-    3:"dynamic degree",
-    4:"motion smoothness",
-    5:'text-to-video alignment',
-    6:'factual consistency',
-}
+aspect_mapping= [
+    "visual quality",
+    "temporal consistency",
+    "dynamic degree",
+    "text-to-video alignment",
+    "factual consistency",
+]
 
 
 @spaces.GPU(duration=60)
@@ -97,11 +96,11 @@ def score(prompt:str, images:List[Image.Image]):
     
     logits = outputs.logits
     num_aspects = logits.shape[-1]
-    aspects = [aspect_mapping[i+1] for i in range(num_aspects)]
+    aspects = [aspect_mapping[i] for i in range(num_aspects)]
     
     aspect_scores = {}
     for i, aspect in enumerate(aspects):
-        aspect_scores[aspect] = logits[0, i].item()
+        aspect_scores[aspect] = round(logits[0, i].item(), 2)
     return aspect_scores