import gradio
import torch
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM

def shorten_text(text, min_length, max_length):
  summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
  short_text = text[:1024]
  summary = summarizer(short_text, max_length, min_length, do_sample=False)
  print("** summary", summary)
  return summary[0]["summary_text"]

def paraphrase_text(text, min_length, max_length):
  tokenizer = AutoTokenizer.from_pretrained("Vamsi/T5_Paraphrase_Paws")  
  model = AutoModelForSeq2SeqLM.from_pretrained("Vamsi/T5_Paraphrase_Paws")
  device = "cuda" if torch.cuda.is_available() else "cpu"
  text_instruction =  "paraphrase: " + text + " </s>"
  encoding = tokenizer.encode_plus(text_instruction, padding="longest", return_tensors="pt")
  input_ids, attention_masks = encoding["input_ids"].to(device), encoding["attention_mask"].to(device)
  outputs = model.generate(
    input_ids=input_ids, attention_mask=attention_masks,
    max_length=max_length,
    do_sample=True,
    top_k=120,
    top_p=0.95,
    early_stopping=True,
    num_return_sequences=5
  )
  line = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
  print("** outputs", len(outputs), line)
  return line

def modify_text(mode, text, min_length, max_length):
  if mode == "shorten":
    return shorten_text(text, min_length, max_length)
  else:
    return paraphrase_text(text, min_length, max_length)

gradio_interface = gradio.Interface(
  fn=modify_text,
  inputs=[
    gradio.Radio(["shorten", "paraphrase"], label="Mode"),
    "text",
    gradio.Slider(5, 200, value=30, label="Min length"),
    gradio.Slider(5, 500, value=130, label="Max length")
  ],
  outputs="text",
  examples=[
    ["shorten",
    """A beautiful golden sun is setting. The sky is on fire. A large neon sign rises into shot. It rests on top of a skyscraper and fills the frame. The building is neither past nor future in design but a bit of both.

Slowly we pan downwards revealing the city that spreads below. A glittering conglomeration of elevated transport tubes, smaller square buildings which are merely huge, with, here and there, the comparatively minuscule relics of previous ages of architecture, pavement level awnings suggesting restaurants and shops. Transparent tubes carry whizzing transport cages past us. An elevated highway carrying traffic composed primarily of large transport lorries passes through frame.

As we descend, the sunlight is blocked out and street lights & neon signs take over as illumination. Eventually we reach the upper levels of a plush shopping precinct.
Xmas decorations are everywhere. People are busy buying, ogling, discussing, choosing wisely from the goodies on display. Shoppers are going by laden with superbly packaged goods. The shop windows are full of elaborately boxed and be-ribboned who-knows-what. In one window is a bank of TV sets on the great majority of the screens is the face of Mr. Helpmann the Deputy Minister of Information. He is being interviewed. No-one bothers to listen to Helpmann.""",
    30, 130]
  ],
  title="Text shortener/paraphraser",
  description="Shortening texts using `facebook/bart-large-cnn`, paraphrasing texts using `Vamsi/T5_Paraphrase_Paws`.",
  article="© Tom Söderlund 2022"
)
gradio_interface.launch()