rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-01-01 02:58:14 +01:00
Code Issues Projects Releases Wiki Activity

Added device selection to transformers pipeline request to utilize GPU for inference

Browse Source
This commit is contained in:
KoboldAI Dev 2021-05-02 23:34:33 -04:00
parent 8dc97ba766
commit 97ad42efe6
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
aiserver.py
View File

@ -108,7 +108,7 @@ if(vars.model != "InferKit"):
print("{0}Initializing transformers, please wait...{1}".format(colors.HEADER, colors.ENDC))
from transformers import pipeline, GPT2Tokenizer
generator = pipeline('text-generation', model=vars.model)
generator = pipeline('text-generation', model=vars.model, device=0)
tokenizer = GPT2Tokenizer.from_pretrained(vars.model)
print("{0}OK! {1} pipeline created!{2}".format(colors.OKGREEN, vars.model, colors.ENDC))
else: