import os
from dotenv import load_dotenv
from modelhub.clients import MLflowClient
from transformers import AutoTokenizer, pipeline

load_dotenv()


def evaluate_model():
    client = MLflowClient(base_url=os.getenv("MODELHUB_BASE_URL"))
    client.set_experiment(os.getenv("MLFLOW_EXPERIMENT_ID"))

    run_id = client.get_previous_stage_run_id()

    test_texts = ["wonderful movie", "terrible film", "great acting", "waste of time"]
    test_labels = [1, 0, 1, 0]

    model_uri = f"runs:/{run_id}/model"
    classifier = client.load_model(model_uri)

    with client.start_run(run_name="evaluate") as run:
        correct = 0
        for text, true_label in zip(test_texts, test_labels):
            pred = classifier(text)
            pred_label = 1 if pred[0]['score'] > 0.5 else 0
            correct += (pred_label == true_label)

        accuracy = correct / len(test_texts)
        client.log_metric("accuracy", accuracy)


if __name__ == "__main__":
    evaluate_model()