import gradio as gr
import os
def func(query):
    try:
        os.system("pip install rwkvstic inquirer transformers torch jax")
        from rwkvstic.load import RWKV
        from rwkvstic.agnostic.backends import JAX
        model = RWKV("https://huggingface.co/BlinkDL/rwkv-4-pile-3b/resolve/main/RWKV-4-Pile-3B-Instruct-test1-20230124.pth",
                    mode=JAX)
        
        def generate(query):
            model.loadContext(newctx=query)
            output = model.forward(number=100)["output"]
            return output
    except Exception as e:
        return "Error [A]" + str(e)
    else:
        try:
            a = generate(query)
            return a
        except Exception as e:
            return "Error [B]" + str(e)

iface = gr.Interface(fn=func, inputs="text", outputs="text")
iface.launch()