File size: 1,033 Bytes
d38c074
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
import huggingface_hub
import os
import argparse
import logging
import torch
from mteb import MTEB
from sentence_transformers import SentenceTransformer
from huggingface_hub import login

try:
    huggingface_hub.login(
        token=os.environ["HUGGINGFACE_TOKEN"], write_permission=True
    )
except Exception as e:
    print(f"Original error: {e}")

def get_device():
    return "cuda" if torch.cuda.is_available() else "cpu"

def load_model(model_id, device):
    model = SentenceTransformer(model_id).to(device)
    logging.info(f"Loaded model {model_id} to the device {device}")
    return model

def run_evaluation(model, model_id,output_folder):
    evaluation = MTEB(task_langs=["es"])
    evaluation.run(model, output_folder=output_folder, eval_splits=["test"])

def evaluate(model_id):
    logging.basicConfig(level=logging.INFO)
    output_folder = f'results/{model_id}'
    device = get_device()
    model = load_model(model_id, device)
    run_evaluation(model, model_id,output_folder)
    
    return output_folder