File size: 847 Bytes
3e56582
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
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()