22from typing import List , Optional
33
44import torch
5- from transformers import AutoAdapterModel , AutoTokenizer
5+ from adapters import AutoAdapterModel
6+ from transformers import AutoTokenizer
67
78logger = logging .getLogger (__name__ )
89
@@ -20,31 +21,46 @@ def build_specter2_text(title: str, abstract: str, tokenizer: AutoTokenizer) ->
2021class Specter2Embedder :
2122 """Wrapper around the SPECTER2 retrieval (proximity) model."""
2223
23- def __init__ (self , model : str = "allenai/specter2" , device : Optional [str ] = None ) -> None :
24+ def __init__ (
25+ self , adapter : str = "allenai/specter2_base" , device : Optional [str ] = None
26+ ) -> None :
2427 """Load SPECTER2 tokenizer, base model, and proximity adapter."""
28+
2529 if device is None :
2630 device = "cuda" if torch .cuda .is_available () else "cpu"
2731
2832 self .device = torch .device (device )
29- logger . info ( "Loading SPECTER2 base + proximity adapter on %s..." , self . device )
33+ self . adapter = adapter
3034
31- # Load tokenizer + base model
32- self .tokenizer = AutoTokenizer .from_pretrained ("allenai/specter2_base" )
35+ # Tokenizer (shared across all variants)
36+ self .tokenizer = AutoTokenizer .from_pretrained (
37+ "allenai/specter2_base" ,
38+ local_files_only = False ,
39+ trust_remote_code = False ,
40+ )
3341
34- self .model = AutoAdapterModel .from_pretrained ("allenai/specter2_base" )
42+ # Adapter-aware model
43+ self .model = AutoAdapterModel .from_pretrained (
44+ "allenai/specter2_base" ,
45+ local_files_only = False ,
46+ trust_remote_code = False ,
47+ )
3548
36- # Load retrieval adapter (proximity)
49+ # Load and activate adapter
3750 self .model .load_adapter (
38- model ,
39- load_as = "specter2" ,
40- set_active = True ,
51+ adapter ,
4152 source = "hf" ,
53+ set_active = True ,
4254 )
4355
4456 self .model .to (self .device )
4557 self .model .eval ()
4658
47- logger .info ("SPECTER2 model + adapter loaded successfully." )
59+ logger .info (
60+ "Loaded SPECTER2 model with adapter '%s' on device '%s'" ,
61+ adapter ,
62+ self .device ,
63+ )
4864
4965 def embed_batch (self , texts : List [str ]) -> List [Optional [List [float ]]]:
5066 """
0 commit comments