import requests 
import json 
import os
import gradio as gr

key = os.getenv('key')
url = os.getenv('entry')
deploy = os.getenv('deploy')
headers = {
    'Content-Type': 'application/json', 
    'Authorization': 'Bearer {}'.format(key), 
    'azureml-model-deployment': deploy,
}

def grin_moe(message, history):
    f_m = ""
    for human, ai in history:
        f_m += '<|user|>\n' + human + '<|end|>\n<|assistant|>\n' + ai + '<|end|>\n'
    f_m += '<|user|>\n' + message + '<|end|>\n<|assistant|>\n'
    
    data = {
        "input_data": {
            "input_string": [f_m],
            "parameters":{
                "max_tokens": 4096, 
                "ignore_eos": False, 
                "stop_token_ids": [32000, 32001, 32007],
                "temperature": 0.0,
            }
        }
    }
    
    response = requests.post(url, json=data, headers=headers)
    response_json = response.json()
    output = response_json[0]["0"][len(f_m):]
    return output 

"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
demo = gr.ChatInterface(
    grin_moe,
    examples=[
        "Sally (a girl) has 3 brothers. Each brother has 2 sisters. How many sisters does Sally have?",
        "How many 'r' in the word 'strawberry'? Go through the count step-by-step",
        "If you pass the second person in a race, what position are you in? Think through the process first",
        "What word is spelled incorrectly in every single dictionary?",
        "If the tangent line to the curve $y = e^x + x$ at the point $(0, 1)$ is also the tangent line to the curve $y = ln(x+1) + a$ at some point, find the value of the constant $a$.",
        "Please complete the following single-choice question. The question has four options, and only one of them is correct. Select the option that is correct. \n\nQuestion: Given $\\cos(\\alpha + \\beta) = m$, and $\\tan(\\alpha) \\tan(\\beta) = 2$, we want to find $\\cos(\\alpha - \\beta) =$\n\nOption A: $-3m$\nOption B: $-\\frac{m}{3}$\nOption C: $\\frac{m}{3}$\nOption D: $3m$"
    ],
    cache_examples=True,
    title="&#128513; MoE",
    description="<font size='+1'>Demo for GRIN MoE (16x3.8B). Full paper to-be-released soon! Note that our model is not very stable on the provided challenging examples yet &#128513;</font><a href='https://huggingface.co/microsoft/GRIN-MoE'>Hugging Face</a> | <a href='https://arxiv.org/pdf/2409.12136'>Paper</a> | <a href='https://github.com/microsoft/GRIN-MoE'>Github</a>"
)


if __name__ == "__main__":
    demo.launch()