mirror of
https://github.com/oobabooga/text-generation-webui.git
synced 2024-12-26 14:20:40 +01:00
107 lines
4.0 KiB
Python
107 lines
4.0 KiB
Python
import os
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Optional, Union
|
|
|
|
import llama_cpp
|
|
import numpy as np
|
|
import torch
|
|
from llama_cpp import Llama
|
|
from torch.nn import CrossEntropyLoss
|
|
from transformers import GenerationConfig, PretrainedConfig, PreTrainedModel
|
|
from transformers.modeling_outputs import CausalLMOutputWithPast
|
|
|
|
from modules import shared
|
|
from modules.llamacpp_model import LlamaCppModel
|
|
from modules.logging_colors import logger
|
|
|
|
|
|
class LlamacppHF(PreTrainedModel):
|
|
def __init__(self, model):
|
|
super().__init__(PretrainedConfig())
|
|
self.model = model
|
|
self.generation_config = GenerationConfig()
|
|
self.cache = None
|
|
|
|
def _validate_model_class(self):
|
|
pass
|
|
|
|
def _validate_model_kwargs(self, model_kwargs: Dict[str, Any]):
|
|
pass
|
|
|
|
def prepare_inputs_for_generation(self, input_ids, **kwargs):
|
|
return {'input_ids': input_ids, **kwargs}
|
|
|
|
@property
|
|
def device(self) -> torch.device:
|
|
return torch.device(0)
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
# TODO: Some decoding methods (such as Contrastive Search) may not work at this time
|
|
assert len(args) == 0, 'no *args should be passed to forward'
|
|
use_cache = kwargs.get('use_cache', True)
|
|
labels = kwargs.get('labels', None)
|
|
seq = kwargs['input_ids'][0].tolist()
|
|
cache = kwargs['past_key_values'] if 'past_key_values' in kwargs else None
|
|
|
|
# Make the forward call
|
|
seq_tensor = torch.tensor(seq)
|
|
self.cache = seq_tensor
|
|
if labels is None:
|
|
if self.cache is None or not torch.equal(self.cache, seq_tensor[:-1]):
|
|
self.model.reset()
|
|
self.model.eval(seq)
|
|
else:
|
|
self.model.eval([seq[-1]])
|
|
|
|
logits = torch.tensor(self.model.eval_logits)[-1].view(1, 1, -1).to(kwargs['input_ids'].device)
|
|
else:
|
|
self.model.reset()
|
|
self.model.eval(seq)
|
|
logits = torch.tensor(self.model.eval_logits)
|
|
logits = logits.view(1, logits.shape[0], logits.shape[1]).to(kwargs['input_ids'].device)
|
|
|
|
# Based on transformers/models/llama/modeling_llama.py
|
|
loss = None
|
|
if labels is not None:
|
|
# Shift so that tokens < n predict n
|
|
shift_logits = logits[..., :-1, :].contiguous()
|
|
shift_labels = labels[..., 1:].contiguous()
|
|
# Flatten the tokens
|
|
loss_fct = CrossEntropyLoss()
|
|
shift_logits = shift_logits.view(-1, logits.shape[-1])
|
|
shift_labels = shift_labels.view(-1)
|
|
# Enable model parallelism
|
|
shift_labels = shift_labels.to(shift_logits.device)
|
|
loss = loss_fct(shift_logits, shift_labels)
|
|
|
|
return CausalLMOutputWithPast(logits=logits, past_key_values=cache if use_cache else None, loss=loss)
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs):
|
|
assert len(model_args) == 0 and len(kwargs) == 0, "extra args is currently not supported"
|
|
if isinstance(pretrained_model_name_or_path, str):
|
|
pretrained_model_name_or_path = Path(pretrained_model_name_or_path)
|
|
|
|
path = Path(f'{shared.args.model_dir}') / Path(pretrained_model_name_or_path)
|
|
if path.is_file():
|
|
model_file = path
|
|
else:
|
|
model_file = list(path.glob('*ggml*.bin'))[0]
|
|
|
|
logger.info(f"llama.cpp weights detected: {model_file}\n")
|
|
params = {
|
|
'model_path': str(model_file),
|
|
'n_ctx': shared.args.n_ctx,
|
|
'seed': int(shared.args.llama_cpp_seed),
|
|
'n_threads': shared.args.threads or None,
|
|
'n_batch': shared.args.n_batch,
|
|
'use_mmap': not shared.args.no_mmap,
|
|
'use_mlock': shared.args.mlock,
|
|
'low_vram': shared.args.low_vram,
|
|
'n_gpu_layers': shared.args.n_gpu_layers,
|
|
'logits_all': True,
|
|
}
|
|
|
|
model = Llama(**params)
|
|
return LlamacppHF(model)
|