Fix HF auth propagation for text encoder loads

app.py

@@ -65,6 +65,9 @@ def _start_text_encoder_server() -> subprocess.Popen:
     hf_token = os.environ.get("HF_TOKEN")
     if hf_token:
         env["HF_TOKEN"] = hf_token
+        env["HUGGING_FACE_HUB_TOKEN"] = hf_token
+        env["HF_HUB_TOKEN"] = hf_token
+        env["HUGGINGFACEHUB_API_TOKEN"] = hf_token
         print(f"[movimento][boot] HF_TOKEN set for text encoder (len={len(hf_token)})")
     else:
         print(f"[movimento][boot] WARNING: HF_TOKEN not found in environment")

kimodo/model/llm2vec/llm2vec.py

@@ -123,12 +123,13 @@ class LLM2Vec(nn.Module):
         # pop out encoder args
         keys = ["pooling_mode", "max_length", "doc_max_length", "skip_instruction"]
         encoder_args = {key: kwargs.pop(key, None) for key in keys if kwargs.get(key) is not None}
+        hf_token = kwargs.get("token")
 
-        tokenizer = AutoTokenizer.from_pretrained(base_model_name_or_path)
+        tokenizer = AutoTokenizer.from_pretrained(base_model_name_or_path, token=hf_token)
         tokenizer.pad_token = tokenizer.eos_token
         tokenizer.padding_side = "left"
 
-        config = AutoConfig.from_pretrained(base_model_name_or_path)
+        config = AutoConfig.from_pretrained(base_model_name_or_path, token=hf_token)
         config_class_name = config.__class__.__name__
 
         model_class = cls._get_model_class(config_class_name, enable_bidirectional=enable_bidirectional)
@@ -146,6 +147,7 @@ class LLM2Vec(nn.Module):
             model = PeftModel.from_pretrained(
                 model,
                 base_model_name_or_path,
+                token=hf_token,
             )
             model = model.merge_and_unload()
 
@@ -153,6 +155,7 @@ class LLM2Vec(nn.Module):
             model = PeftModel.from_pretrained(
                 model,
                 peft_model_name_or_path,
+                token=hf_token,
             )
             if merge_peft:
                 model = model.merge_and_unload()

kimodo/model/llm2vec/llm2vec_wrapper.py

@@ -24,6 +24,7 @@ class LLM2VecEncoder:
         self.llm_dim = llm_dim
 
         cache_dir = os.environ.get("HUGGINGFACE_CACHE_DIR")
+        hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
 
         if "TEXT_ENCODERS_DIR" in os.environ:
             base_model_name_or_path = os.path.join(os.environ["TEXT_ENCODERS_DIR"], base_model_name_or_path)
@@ -34,6 +35,7 @@ class LLM2VecEncoder:
             peft_model_name_or_path=peft_model_name_or_path,
             torch_dtype=torch_dtype,
             cache_dir=cache_dir,
+            token=hf_token,
         )
         self.model.eval()
         for p in self.model.parameters():

kimodo/scripts/run_text_encoder_server.py

@@ -4,6 +4,8 @@
 import argparse
 import os
 
+os.environ.pop("GRADIO_HOT_RELOAD", None)
+
 import gradio as gr
 import numpy as np
 
@@ -125,7 +127,7 @@ def main():
     # Model will be loaded lazily on first request
     demo_wrapper_fn = DemoWrapper(args.text_encoder, args.tmp_folder)
 
-    with gr.Blocks(title="Text encoder", css=css, theme=theme) as demo:
+    with gr.Blocks(title="Text encoder") as demo:
         gr.Markdown(f"# Text encoder: {display_name}")
         gr.Markdown("## Description")
         gr.Markdown("Get a embeddings from a text.")
@@ -190,7 +192,7 @@ def main():
         )
         clear.click(fn=clear_fn, inputs=None, outputs=outputs)
 
-    demo.launch(server_name=server_name, server_port=server_port)
+    demo.launch(server_name=server_name, server_port=server_port, theme=theme, css=css)
 
 
 if __name__ == "__main__":