BART model small/main.py

from transformers import pipeline, BartForConditionalGeneration, BartTokenizer


# Define the directory where you want to save the tokenizer files
tokenizer_directory = '/Users/vishnu/Downloads/BART model small/model'

# Download and save the tokenizer files from the original BART model
tokenizer = BartTokenizer.from_pretrained('facebook/bart-large-cnn')
tokenizer.save_pretrained(tokenizer_directory)

print(f"Tokenizer saved to {tokenizer_directory}")

# Load the fine-tuned model and tokenizer
model_directory = './model'
tokenizer = BartTokenizer.from_pretrained(model_directory)
model = BartForConditionalGeneration.from_pretrained(model_directory)

# Create a summarization pipeline
summarizer = pipeline('summarization', model=model, tokenizer=tokenizer)

# Sample news article
article = """
NASA's James Webb Space Telescope has captured the most detailed images of Jupiter's atmosphere ever seen. 
The stunning visuals showcase the planet's turbulent storms, swirling winds, and vibrant auroras, offering 
new insights into its complex weather patterns. These observations were made possible thanks to Webb's powerful 
infrared capabilities, which allow scientists to study the gas giant in unprecedented detail. The telescope's 
latest findings could shed light on how Jupiter's atmosphere and weather systems function, helping us understand 
similar processes on other planets.
"""

# Generate the summary
summary = summarizer(article, max_length=128, min_length=30, length_penalty=2.0, num_beams=4, early_stopping=True)

# Output the summary
print("\n\n Summary:", summary[0]['summary_text'])