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 17:34:48 +08:00 • committed by GitHub
parent 9285750c2a
commit fea00a66fd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 5 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

10
wfgy_core.py
View file

@ -4,16 +4,16 @@ from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from huggingface_hub import InferenceClient
class WFGYRunner:
def __init__(self, model_id="HuggingFaceH4/zephyr-7b-alpha", use_remote=True):
def __init__(self, model_id="tiiuae/falcon-7b-instruct", use_remote=True):
self.use_remote = use_remote
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.model_id = model_id
if self.use_remote:
try:
self.client = InferenceClient(model=self.model_id, token=os.environ.get("HF_TOKEN"))
except Exception as e:
raise RuntimeError(f"Hugging Face remote mode failed: {e}")
token = os.environ.get("HF_TOKEN")
if not token:
raise RuntimeError("Missing HF_TOKEN environment variable.")
self.client = InferenceClient(model=self.model_id, token=token)
else:
self.tokenizer = AutoTokenizer.from_pretrained(model_id)
self.model = AutoModelForCausalLM.from_pretrained(model_id)