Spaces:
Sleeping
Sleeping
Update src/macg/llm_hf.py
Browse files- 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 |
-
|
| 31 |
-
|
| 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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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": {
|