import gradio as gr
import os
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
import json

access_token = os.environ['ACCES_TOKEN']

model = AutoModelForSequenceClassification.from_pretrained("EkhiAzur/C1_Sailkapen_Demoa", token=access_token)

tokenizer = AutoTokenizer.from_pretrained(
  "EkhiAzur/C1_Sailkapen_Demoa",
  token = access_token,
  use_fast=True,
  add_prefix_space=True,
)

classifier = pipeline("text-classification", tokenizer=tokenizer, model=model, max_length=512,
                padding=True, truncation=True, batch_size=1)

adibideak = json.load(open("./Adibideak.json", "r"))

def prozesatu(Testua, request: gr.Request):
    if Testua[-3:]=="...":
        Testua = prozesatu.adibideak[Testua]
    prediction = prozesatu.classifier(Testua)[0]
    if prediction["label"]=="GAI":
        return {"Gai":prediction["score"], "Ez gai": 1-prediction["score"]}
    else:
        return {"Gai":1-prediction["score"], "Ez gai": prediction["score"]}

def testua_itzuli(testua):
    if testua not in testua_itzuli.adibideak:
        return ""
    return testua_itzuli.adibideak[testua]
    
testua_itzuli.adibideak = adibideak
prozesatu.adibideak = adibideak
prozesatu.classifier = classifier

def ezabatu(Testua):
  return ""

with gr.Blocks() as demo:
  with gr.Row():
    with gr.Column():
      input = gr.Textbox(label="Testua")
      with gr.Row():
        bidali_btn = gr.Button("Bidali")
        ezabatu_btn = gr.Button("Ezabatu")
    
    label = gr.Label(num_top_classes=2, label="C1 maila")
  
  bidali_btn.click(fn=prozesatu, inputs=input, outputs=label)
  ezabatu_btn.click(fn=ezabatu, inputs=input, outputs=input)
    
  gr.Examples(list(adibideak.keys()), inputs=input, outputs=input, label="Adibideak:", fn=testua_itzuli, cache_examples=True)

demo.launch()