chatbot.py

# from rwkvstic.load import RWKV
# from rwkvstic.agnostic.backends import TORCH
# import torch
# import os

# os.environ["RWKV_JIT_ON"] = '1'
# os.environ["RWKV_CUDA_ON"] = '1' # if '1' then use CUDA kernel for seq mode (much faster)
# # this is the dtype used for trivial operations, such as vector->vector operations and is the dtype that will determine the accuracy of the model
# runtimedtype = torch.float32 # torch.float64, torch.bfloat16

# # this is the dtype used for matrix-vector operations, and is the dtype that will determine the performance and memory usage of the model
# dtype = torch.bfloat16 # torch.float32, torch.float64, torch.bfloat16

# useGPU = False # False

# model = RWKV("RWKV-4-Pile-3B-Instruct-test2-20230209.pth", mode=TORCH, useGPU=useGPU, runtimedtype=runtimedtype, dtype=dtype)

# model.loadContext(newctx=f"Q: How many hours are there in a day?\n\nA:")
# output = model.forward(number=100)["output"]

# print(output)