from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

tokenizer = AutoTokenizer.from_pretrained("rristo/mlong-t5-tglobal-base-et-riigikogu-summary")
model = AutoModelForSeq2SeqLM.from_pretrained("rristo/mlong-t5-tglobal-base-et-riigikogu-summary")

text="""Varasematest uuringutest on teada, et punetav nägu võib märku anda erutusest näiteks aaradel ja raisakotkastel. Sestap huvitas Tours'i Ülikooli etoloog Delphine Soulet'd ja tema kolleege, kas sarnast tundemärki võib näha ka kodukanade (Gallus gallus domesticus) nägudel.
Töörühm filmis esmalt kuut Sussexi tõugu kana erinevates olukordades. Mõnes olukorras toimetasid kanad loomulikult omasoodu, teistes aga juhtisid uurijad lindude tegevust. Põnevates ja autasu tõotavates olukordades lasi töörühm kanadel võtta tolmuvanni või söötis neid ussikestega. Hirmuga seotud olukordades püüdsid uurijad linde käsitsi kinni.
Katsete järel oli töörühma päralt videosalvestistest võetud tuhandeid üksikkaadreid. Just nende analüüsiks loodud algoritmi toel said uurijad täpselt jälgida, kui punased olid igas olukorras kanade hari, põsed, kõrvanibud ja lotid.
Töörühma sõnul oli uuringu valim väike, mistõttu vajavad tulemused kinnitamist suuremas kordusuuringus. Siiski ilmneb tulemustest, et vähem punetavad põsed ja kõrvanibud võivad viidata linnu rahulikule ja rõõmsale seisundile. Vastukaaluks näib punetavam nägu märku andvat linnu suuremast emotsionaalsest erutusest. Sinna hulka kuuluvad nii ussikeste saamisega seotud elevus kui ka hirm.
Soulet ja kolleegid tegid veel ühe katse, kus jaotasid 25 Sussexi tõugu kana kahte rühma. Uurijad käisid viie nädala jooksul 13 linnu juures, et kanu pisitasa inimese kohaoluga harjutada. Ülejäänud 12 lindu jäeti viieks nädalaks kontrollrühmana omapäi.
Kui siis kõik kanad viie nädala möödudes uuesti inimestega kokku puutusid, ilmnes kahe kanarühma vahel selge vahe. Uurijatega harjunud linnud pelgasid inimest vähem ja muutusid nende juuresolekul näost vähem punaseks, kui nende üksi jäetud liigikaaslased."""


def summarize(text, model, tokenizer, max_new_tokens=512, device='cuda'):
    input_ids = tokenizer(
         text, return_tensors="pt"
    ).input_ids  # Batch size 1
    outputs = model.generate(input_ids=input_ids.to(device), max_new_tokens=max_new_tokens)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

DEVICE='cuda'
model=model.to(DEVICE)

summarize(text, model, tokenizer, device=DEVICE)