from transformers import pipeline
import gradio as gr
import spaces
import os
from threading import Thread
from typing import Iterator
import torch

from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    TextIteratorStreamer
)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
tokenizer = AutoTokenizer.from_pretrained("dascim/greekbart-news24-abstract")
model = AutoModelForSeq2SeqLM.from_pretrained("dascim/greekbart-news24-abstract")
model.eval()
@spaces.GPU(duration=90)
def get_input(text) -> Iterator[str]:
  streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_special_tokens=True)
  input_ids = tokenizer.encode(text, add_special_tokens=True, return_tensors='pt')
  
  generate_kwargs = dict(
        input_ids=input_ids,
        tokenizer=tokenizer,
        streamer=streamer,
        max_new_tokens=120,
        do_sample=False,
        num_beams=1,
        early_stopping=False
    )

  t = Thread(target=model.generate, kwargs=generate_kwargs)
  t.start()

  outputs = []
  for text in streamer:
      outputs.append(text)
      yield "".join(outputs)
  
desc = f'''
This is a demo for Greek News summarization using [greekbart-news24-abstract](https://huggingface.co/dascim/greekbart-news24-abstract), a finetuned version of [GreekBART](https://huggingface.co/dascim/greekbart). 
GreekBART is the first Greek sequence to sequence pretrained model. It is pretrained on 77GB of Greek raw text using the CNRS Jean Zay supercomputer. Our model is based on BART. Unlike already existing BERT-based Greek language models such as GreekBERT and Electra, GreekBART is particularly well-suited for generative tasks, since not only its encoder but also its decoder is pretrained.

📑 Paper: [GreekBART: The First Pretrained Greek Sequence-to-Sequence Model](https://arxiv.org/abs/2304.00869)

Enter your text (maximum of 1024 tokens of Greek news article) to get a summary.
'''
iface = gr.Interface(fn=get_input,inputs="text",outputs="text",title = "🇬🇷 Greek News Summarizer 🇬🇷",description=desc, examples=[['Στην κατάθεση νοσηλεύτριας του Καραμανδάνειου Νοσοκομείου Πάτρας Παναγιώτας Τσεντούρου, η οποία εργαζόταν όταν εισήχθη στις 8 Απριλίου 2021 η Τζωρτζίνα, προχώρησε η διαδικασία ενώπιον του ΜΟΔ που δικάζει τη Ρούλα Πισπιρίγκου. Η νοσηλεύτρια κατέθεσε πως κατά την εισαγωγή του παιδιού "μου ανέφεραν πως είναι ένα παιδάκι που έχει χάσει τα αδελφάκια του και ότι είναι ιδιαίτερη περίπτωση” και εξιστόρησε τα γεγονότα της ημέρας εισαγωγής και της επομένης που η ίδια είχε βάρδια στην παιδιατρική κλινική.']])
iface.launch()