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/RoBERTA_3", token=access_token)

tokenizer = AutoTokenizer.from_pretrained(
  "EkhiAzur/RoBERTA_3",
  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:-1]=="...":
       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"]}

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(['Azken urte hauetan, japonen telebista oso esaguna da...', 'Denok bizi izan genuen iaz ia hiru hilabetez jardun zen...', 'Pandemia egoera eta konfinamendua dela eta, asko izan dira...'], inputs=input, outputs=label, label="Adibideak:", fn=prozesatu, cache_examples=True)

demo.launch()