# 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)