File size: 1,442 Bytes
c020db5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
from typing import Dict, List
from transformers import AutoTokenizer
from adapters import AutoAdapterModel
class AdapterHandler:
def __init__(self):
self.tokenizer = None
self.model = None
def initialize(self, model_dir: str):
self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
self.model = AutoAdapterModel.from_pretrained(model_dir)
# Load adapter - adjust path as needed
self.model.load_adapter("specter2_proximity", source="local")
self.model.set_active_adapters("specter2_proximity")
return self
def __call__(self, data: Dict[str, List[str]]) -> Dict[str, List[float]]:
titles = data.get("title", [""])
abstracts = data.get("abstract", [""] * len(titles))
# Combine inputs
combined = [
title + self.tokenizer.sep_token + (abstract or "")
for title, abstract in zip(titles, abstracts)
]
# Tokenize
inputs = self.tokenizer(
combined,
padding=True,
truncation=True,
return_tensors="pt",
return_token_type_ids=False,
max_length=512
)
# Get embeddings
with torch.no_grad():
outputs = self.model(**inputs)
embeddings = outputs.last_hidden_state[:, 0, :].numpy()
return {"embeddings": embeddings.tolist()}
|