text-generation-webui/modules/RWKV.py

46 lines
1.4 KiB
Python
Raw Normal View History

2023-02-28 04:09:11 +01:00
import os
import time
import types
2023-02-28 03:50:16 +01:00
from pathlib import Path
2023-02-28 04:09:11 +01:00
2023-02-28 03:50:16 +01:00
import numpy as np
2023-02-28 04:09:11 +01:00
import torch
import modules.shared as shared
2023-02-28 03:50:16 +01:00
np.set_printoptions(precision=4, suppress=True, linewidth=200)
os.environ['RWKV_JIT_ON'] = '1'
os.environ["RWKV_CUDA_ON"] = '0' # '1' : use CUDA kernel for seq mode (much faster)
from rwkv.model import RWKV
from rwkv.utils import PIPELINE, PIPELINE_ARGS
2023-03-01 16:18:17 +01:00
2023-03-01 16:08:55 +01:00
class RWKVModel:
def __init__(self):
pass
2023-02-28 03:50:16 +01:00
2023-03-01 16:08:55 +01:00
@classmethod
def from_pretrained(self, path, dtype="fp16", device="cuda"):
tokenizer_path = Path(f"{path.parent}/20B_tokenizer.json")
2023-02-28 03:50:16 +01:00
2023-03-01 16:08:55 +01:00
model = RWKV(model=path.as_posix(), strategy=f'{device} {dtype}')
pipeline = PIPELINE(model, tokenizer_path.as_posix())
2023-02-28 03:50:16 +01:00
2023-03-01 16:08:55 +01:00
result = self()
result.model = pipeline
return result
2023-03-01 16:16:11 +01:00
def generate(self, context, token_count=20, temperature=1, top_p=1, alpha_frequency=0.25, alpha_presence=0.25, token_ban=[0], token_stop=[], callback=None):
args = PIPELINE_ARGS(
temperature = temperature,
top_p = top_p,
2023-03-01 16:19:37 +01:00
alpha_frequency = alpha_frequency, # Frequency Penalty (as in GPT-3)
alpha_presence = alpha_presence, # Presence Penalty (as in GPT-3)
token_ban = token_ban, # ban the generation of some tokens
token_stop = token_stop
2023-03-01 16:16:11 +01:00
)
return self.model.generate(context, token_count=token_count, args=args, callback=callback)