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

load_dotenv()

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

    texts = ["good", "bad", "excellent", "poor"]
    labels = [1, 0, 1, 0]

    model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)
    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
    classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)

    with client.start_run(run_name="simple_train") as run:
        for epoch in range(2):
            for text, label in zip(texts, labels):
                outputs = classifier(text)
                client.log_metric("prediction", outputs[0]['score'])

        client.mlflow.transformers.log_model(classifier, "simple_model")

if __name__ == "__main__":
    train_model()