import gradio as gr
import os
def func(query):
    try:
        os.system("pip install rwkvstic inquirer transformers torch jax jaxlib")
        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()