mmarch/text-generation-webui
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Add CUDA wheels for llama-cpp-python by jllllll

Browse source
This commit is contained in:
oobabooga 2023-07-19 19:31:19 -07:00
parent 05f4cc63c8
commit 4b19b74e6c
3 changed files with 19 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

7
modules/llamacpp_model.py
View file

@ -9,12 +9,17 @@ https://abetlen.github.io/llama-cpp-python/
import re
from functools import partial
from llama_cpp import Llama, LlamaCache, LogitsProcessorList
import torch
from modules import shared
from modules.callbacks import Iteratorize
from modules.logging_colors import logger
if torch.cuda.is_available():
from llama_cpp_cuda import Llama, LlamaCache, LogitsProcessorList
else:
from llama_cpp import Llama, LlamaCache, LogitsProcessorList
def ban_eos_logits_processor(eos_token, input_ids, logits):
logits[eos_token] = -float('inf')