remiai3's picture
Upload 4 files
3e56582 verified
import argparse, torch
from transformers import pipeline
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--text", type=str, default=(
"Artificial intelligence is transforming education by enabling personalized learning. "
"Teachers can use AI-driven tools to understand student progress and tailor activities."
))
parser.add_argument("--max_length", type=int, default=120)
parser.add_argument("--min_length", type=int, default=40)
args = parser.parse_args()
device = 0 if torch.cuda.is_available() else -1
summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device=device)
out = summarizer(args.text, max_length=args.max_length, min_length=args.min_length, do_sample=False)
print(out[0]["summary_text"])
if __name__ == "__main__":
main()