The open engine for
deep learning
workflows.

# fenn.yaml
project: lora-seq-cls

general:
  device: cuda

model:
  name: distilbert-base-uncased
  num_labels: 2
  max_length: 128

train:
  seed: 42
  epochs: 3
  lr: 2e-4
  batch: 16

lora:
  r: 8
  alpha: 16
  dropout: 0.1

from fenn import Fenn
from fenn.nn.trainers import LoRATrainer
import torch.optim as optim
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from datasets import load_dataset
from sklearn.metrics import accuracy_score

app = Fenn()

@app.entrypoint
def main(args):
    raw = load_dataset("sst2")
    tokenizer = AutoTokenizer.from_pretrained(args["model"]["name"])
    # ... build SentimentDataset and DataLoaders ...

    model = AutoModelForSequenceClassification.from_pretrained(
        args["model"]["name"], num_labels=args["model"]["num_labels"]
    )
    optimizer = optim.AdamW(
        model.parameters(), lr=float(args["train"]["lr"])
    )
    trainer = LoRATrainer(
        model=model, optim=optimizer, task_type="SEQ_CLS",
        r=args["lora"]["r"], lora_alpha=args["lora"]["alpha"],
        lora_dropout=float(args["lora"]["dropout"]),
        target_modules=["q_lin", "v_lin"],
        device=args["general"]["device"],
    )
    trainer.fit(train_loader, epochs=args["train"]["epochs"])
    predictions = trainer.predict(test_loader)
    print(f"Accuracy: {accuracy_score(test_labels, predictions):.4f}")

if __name__ == "__main__":
    app.run()