package io.synor.compute;

/**
 * Options for attention operations.
 */
public class AttentionOptions {
    private final int numHeads;
    private final boolean flash;
    private final Precision precision;
    private final ProcessorType processor;

    private AttentionOptions(Builder builder) {
        this.numHeads = builder.numHeads;
        this.flash = builder.flash;
        this.precision = builder.precision;
        this.processor = builder.processor;
    }

    public int getNumHeads() {
        return numHeads;
    }

    public boolean isFlash() {
        return flash;
    }

    public Precision getPrecision() {
        return precision;
    }

    public ProcessorType getProcessor() {
        return processor;
    }

    public static Builder builder() {
        return new Builder();
    }

    public static class Builder {
        private int numHeads = 8;
        private boolean flash = true;
        private Precision precision = Precision.FP16;
        private ProcessorType processor = ProcessorType.GPU;

        public Builder numHeads(int numHeads) {
            this.numHeads = numHeads;
            return this;
        }

        public Builder flash(boolean flash) {
            this.flash = flash;
            return this;
        }

        public Builder precision(Precision precision) {
            this.precision = precision;
            return this;
        }

        public Builder processor(ProcessorType processor) {
            this.processor = processor;
            return this;
        }

        public AttentionOptions build() {
            return new AttentionOptions(this);
        }
    }
}
