import os
API_TOKEN = os.getenv('API_TOKEN')

import gradio as gr
import requests

API_URL = "https://api-inference.huggingface.co/models/tlkh/flan-t5-paraphrase-classify-explain"
headers = {"Authorization": f"Bearer {API_TOKEN}"}

def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.json()

def infer(s1, s2):
    model_input = "Classify and explain the relationship between this pair of sentences: <S1> "+s1+" </S1><S2> "+s2+" </S2>"
    data = query(model_input)

    if "error" in data:
        return "Error: "+ data["error"]
    elif "generated_text" in data[0]:
        return data[0]["generated_text"]
    else:
        return data

title = "Paraphrase Classification and Explanation"
desc = "Classify and explain the semantic relationship between the two sentences"
long_desc = "This is a Flan-T5-Large model fine-tuned to perform paraphrase classification and explanation. It takes in two sentences as inputs."

s1 = gr.Textbox(value="On Monday, Tom went to the market.",label="Sentence 1")
s2 = gr.Textbox(value="Tom went to the market.",label="Sentence 2")

examples = gr.Examples(
        examples=[
            ["one", "two"],
        ],
        inputs=[s1, s2],
)

demo = gr.Interface(fn=infer, inputs=[s1,s2], outputs="text",
                    title=title,
                    description=desc,
                    article=long_desc,
                    )

demo.launch()