Desiring Retriever Only Inference about dalm HOT 2 CLOSED

from dalm.models.retriever_only_base_model import AutoModelForSentenceEmbedding

model = "arcee-ai/bge-code-retriever"

retriever_model = AutoModelForSentenceEmbedding(
        model,
        get_peft=False,
        use_bnb=False,
        is_autoregressive=True,
    )

tokenizer = retriever_model.tokenizer

passage = """class RetrieverModel:
    def __init__(self, base_model=DEFAULT_BASE_MODEL):
        self.tokenizer = AutoTokenizer.from_pretrained(base_model)
        self.model = SentenceTransformer(base_model)
        self.base_model = base_model
        self.peft_model = None
        self.on_cuda = False
        self.device = None
        # special code retriever model
        self.is_code_retriever = False

    def to(self, device):
        self.model.to(device)
        self.on_cuda = True
        self.device = device"""

model_inputs = tokenizer(passage, return_tensors="pt")

model_inputs.to(retriever_model.device)

embeddings = retriever_model(model_inputs["input_ids"], model_inputs["attention_mask"])

print(embeddings)