from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
import gradio as gr

# Load model directly

tokenizer = AutoTokenizer.from_pretrained("ahmedmbutt/PTS-Bart-Large-CNN")
model = AutoModelForSeq2SeqLM.from_pretrained("ahmedmbutt/PTS-Bart-Large-CNN")

def summarize (text):

    input_ids = tokenizer(text, return_tensors="pt").input_ids
    
    #run through model - use [0] to select the tensor as .generate returns a list
    output_ids = model.generate(
    input_ids,
    max_new_tokens=150, 
    no_repeat_ngram_size=4, 
    do_sample=True
    )[0]
    
    output = tokenizer.decode (output_ids, skip_special_tokens=True)
    
    return (output)



iface = gr.Interface(fn=summarize, 
                     inputs=[gr.Textbox(label="Paste text to summarize and hit 'Submit'.", lines=10)],
                     outputs=[gr.Textbox(label="Summarized text:", lines=10)],
                     title="AI Text Summarizer", 
                     #description="Turn keywords into creative suggestions",
                     article="""
                     <div align=left>
                     <h1>AI Text Summarizer</h1>
                     <li>Generate a succinct and overall summary of text of any length.</li>
                     <li>Many AI models fail to generate a good summary - they either pick up first sentences, or fail to blend information together</li>
                     <li>The AI text summarizer is well trained  to create better summaries that incorporate details and put facts together.</li>
                     <li>Click on the example for a demo</li>

                     <p>AI Model:<br>
                     <li>Bart-Large-CNN trained on the PTS dataset.</li>
                     <li>Original model id: ahmedmbutt/PTS-Bart-Large-CNN</li>
                     
                     </p><p>Default parameter details:<br>
                     <li><code>do_sample=True</code>, <code>no_repeat_ngram_size=3</code>,<code>max_new_tokens = 30</code></li>
                     </div>
                     
                     """,
                     # flagging_mode='never',
                     examples= ["""What is driving food inflation higher? 
A drought last year and an ongoing heat wave have significantly reduced the supplies of foods like pulses, vegetables, and cereals. Curbs on food exports and reducing tariffs on imports have had little effect. Although vegetable supplies generally decrease during the summer months, this year's decline is much more pronounced. Temperatures in nearly half of the country are soaring 4-9 degrees Celsius above normal, spoiling harvested and stored vegetables and hindering the planting of crops such as onions, tomatoes, eggplant and spinach."""],  
                     cache_examples=True,
                    )

iface.launch()


# In[ ]:

# #### Run Using Inference Endpoing  - Serverless with Access Token generated on HF

# In[ ]:


#import requests

#API_URL = "https://api-inference.huggingface.co/models/facebook/bart-large-cnn"
#headers = {"Authorization": 'Bearer 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'}

#def query(payload):
    #response = requests.post(API_URL, headers=headers, json=payload)
    #return response.json()

#output = query({
#"inputs": "The tower is 324 metres (1,063 ft) tall, about the same height as an 81-storey building, and the tallest structure in Paris. Its base is square, measuring 125 metres (410 ft) on each side. During its construction, the Eiffel Tower surpassed the Washington Monument to become the tallest man-made structure in the world, a title it held for 41 years until the Chrysler Building in New York City was finished in 1930. It was the first structure to reach a height of 300 metres. Due to the addition of a broadcasting aerial at the top of the tower in 1957, it is now taller than the Chrysler Building by 5.2 metres (17 ft). Excluding transmitters, the Eiffel Tower is the second tallest free-standing structure in France after the Millau Viaduct.",
#})