vrr/WFGY
2
0
Fork 0
mirror of https://github.com/onestardao/WFGY.git synced 2026-04-28 03:29:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

Update wfgy_core.py

Browse source
This commit is contained in:
PSBigBig 2025-06-11 16:16:27 +08:00 • committed by GitHub
parent 6b06c35eec
commit be1036076c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 4 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

7
wfgy_core.py
View file

@ -2,14 +2,15 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from huggingface_hub import InferenceClient
import os
class WFGYRunner:
def __init__(self, model_id="mistralai/Mixtral-8x7B-Instruct-v0.1", hf_token=None, use_remote=False):
def __init__(self, model_id="tiiuae/falcon-7b-instruct", hf_token=None, use_remote=False):
self.model_id = model_id
self.hf_token = hf_token or os.getenv("HF_TOKEN")
self.use_remote = use_remote
self.hf_token = hf_token or os.environ.get("HF_TOKEN")
if not self.hf_token:
raise ValueError("HF_TOKEN not found in environment variables. Set it before running.")
raise ValueError("HF_TOKEN not found in environment variables or passed to constructor.")
if self.use_remote:
self.client = InferenceClient(model=model_id, token=self.hf_token)