oobabooga/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2026-04-20 14:03:45 +00:00
Code Issues Projects Releases Packages Wiki Activity

Make llama-cpp-python not crash immediately

Browse source
This commit is contained in:
oobabooga 2024-07-04 19:15:37 -07:00
parent f77cf159ba
commit f243b4ca9c
3 changed files with 37 additions and 73 deletions
Download patch file Download diff file Expand all files Collapse all files

29
modules/llamacpp_model.py
View file

@ -4,37 +4,12 @@ from functools import partial
import numpy as np
import torch
from modules import llama_cpp_python_hijack, shared
from modules import shared
from modules.callbacks import Iteratorize
from modules.llama_cpp_python_hijack import llama_cpp_lib
from modules.logging_colors import logger
from modules.text_generation import get_max_prompt_length
try:
import llama_cpp
except:
llama_cpp = None
try:
import llama_cpp_cuda
except:
llama_cpp_cuda = None
try:
import llama_cpp_cuda_tensorcores
except:
llama_cpp_cuda_tensorcores = None
def llama_cpp_lib():
if shared.args.cpu and llama_cpp is not None:
return llama_cpp
elif shared.args.tensorcores and llama_cpp_cuda_tensorcores is not None:
return llama_cpp_cuda_tensorcores
elif llama_cpp_cuda is not None:
return llama_cpp_cuda
else:
return llama_cpp
def ban_eos_logits_processor(eos_token, input_ids, logits):
logits[eos_token] = -float('inf')