import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

model_path = "anilbhatt1/phi2-oasst-guanaco-bf16-custom"
model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_path)

def generate_text(prompt, response_length):   
            
    prompt = str(prompt)
    max_len = int(response_length)    

    gen = pipeline('text-generation', model=model, tokenizer=tokenizer, max_length=max_len)
    result = gen(prompt)
    output_msg = result[0]['generated_text']
    return output_msg

def gradio_fn(prompt, response_length):
    output_txt_msg = generate_text(prompt, response_length)    
    return output_txt_msg

markdown_description = """
- This is a Gradio app that answers the query you ask it
- Uses **microsoft/phi-2 qlora** optimized model finetuned on **timdettmers/openassistant-guanaco** dataset 
"""
demo = gr.Interface(fn=gradio_fn,
                    inputs=[gr.Textbox(info="How may I help you ? please enter your prompt here..."),
                            gr.Slider(value=50, minimum=50, maximum=200, \
                                      info="Choose a response length min chars=50, max=200")],                       
                    outputs=gr.Textbox(),
                    title="phi2 - Dialog Partner",
                    description=markdown_description,
                    article=" **Credits** : https://github.com/mshumer/gpt-llm-trainer ")

demo.queue().launch(share=True)