Navya-Sree commited on
Commit
ae9979f
·
verified ·
1 Parent(s): b1af075

Update src/macg/llm_hf.py

Browse files
Files changed (1) hide show
  1. src/macg/llm_hf.py +8 -4
src/macg/llm_hf.py CHANGED
@@ -27,9 +27,8 @@ class HuggingFaceInferenceLLM(LLMClient):
27
  timeout_s: int = 90,
28
  ) -> None:
29
  self.model = model
30
- self.token = token or os.getenv("HF_TOKEN")
31
- if not self.token:
32
- raise ValueError("HF_TOKEN is not set.")
33
  self.max_new_tokens = max_new_tokens
34
  self.temperature = temperature
35
  self.retries = retries
@@ -37,7 +36,12 @@ class HuggingFaceInferenceLLM(LLMClient):
37
 
38
  def complete(self, system: str, prompt: str) -> str:
39
  url = f"https://api-inference.huggingface.co/models/{self.model}"
40
- headers = {"Authorization": f"Bearer {self.token}"}
 
 
 
 
 
41
  payload = {
42
  "inputs": f"{system}\n\n{prompt}".strip(),
43
  "parameters": {
 
27
  timeout_s: int = 90,
28
  ) -> None:
29
  self.model = model
30
+ # allow either HF_TOKEN or HUGGINGFACEHUB_API_TOKEN, and allow missing token
31
+ self.token = token or os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
 
32
  self.max_new_tokens = max_new_tokens
33
  self.temperature = temperature
34
  self.retries = retries
 
36
 
37
  def complete(self, system: str, prompt: str) -> str:
38
  url = f"https://api-inference.huggingface.co/models/{self.model}"
39
+
40
+ # ✅ if token exists, use it; otherwise call unauthenticated
41
+ headers = {}
42
+ if self.token:
43
+ headers = {"Authorization": f"Bearer {self.token}"}
44
+
45
  payload = {
46
  "inputs": f"{system}\n\n{prompt}".strip(),
47
  "parameters": {