import gradio as gr
from llm_rs import AutoModel,SessionConfig,GenerationConfig,Precision

repo_name = "rustformers/mpt-7b-ggml"
file_name = "mpt-7b-instruct-q5_1-ggjt.bin"

session_config = SessionConfig(threads=2,batch_size=2)
model = AutoModel.from_pretrained(repo_name, model_file=file_name, session_config=session_config,verbose=True)

def process_stream(instruction):

    prompt=f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.
### Instruction:
{instruction}
### Response:
Answer:"""
    generation_config = GenerationConfig(seed=40,temperature=0.1,top_p=0.15,top_k=20,max_new_tokens=1000)
    response = ""
    streamer = model.stream(prompt=prompt,generation_config=generation_config)
    for new_text in streamer:
        response += new_text
        yield response
inputs = gr.inputs.Textbox(lines=7, label="Chat with AI")
outputs = gr.outputs.Textbox(label="Reply")

gr.Interface(fn=process_stream, inputs=inputs, outputs=outputs, title="Self_Trained_V2",
             description="Ask anything you want",
             ).launch()