import gradio as gr
import markdownify
import html as pyhtml
from setfit import SetFitModel

model = SetFitModel.from_pretrained("./predictor_2")


def clean_text(text):
  text = markdownify.markdownify(
      pyhtml.unescape(text or ""),
      heading_style="ATX").strip()
  if len(text) > 250:
    text = text[:250] + "..."
  return text

def text_template(title="",description="",url="",comment=""):
  description = clean_text(description)
  comment = clean_text(comment)
  return f"""Title: "{title}"
Description: "{description}"
Url: "{url}"
First Comment: "{
      comment
      }\""""

def text_classifier(title="",description="",url="",comment=""):
    text = text_template(title,description,url,comment)
    prediction = model.predict_proba([text])[0]
    a,b = prediction
    return {'non-ai': float(a), 'ai': float(b)}
inputs = [
                        gr.Textbox(
                            value="",
                            label="Title"
                                   ),
                        gr.Textbox(
                            value="",
                            label="Description"
                                   ),
                        gr.Textbox(
                            value="",
                            label="URL"
                                   ),
                        gr.Textbox(
                            value="",
                            label="Comment"
                                   ),
                        ]



demo = gr.Interface(fn=text_classifier,
                    inputs=inputs,
                    outputs="label")
demo.launch(show_api=True)