Views
No views yet
facebook/bart-base) on custom summarization tasks. After training, the model can generate summaries for input text, which can be used for various applications like news article summarization, report generation, etc.1pip install transformers torch huggingface_hub
2Usage
3Loading the Model and Tokenizer
4Ensure you have saved your trained model and tokenizer in the ./custom_bart_model directory. The code snippet below demonstrates how to load the model and generate summaries based on user input.
5
6from transformers import BartTokenizer, BartForConditionalGeneration
7import torch
8
9# Load the model and tokenizer
10model = "rohansb10/summary"
11tokenizer = BartTokenizer.from_pretrained(model)
12model = BartForConditionalGeneration.from_pretrained(model)
13
14# Move model to the appropriate device
15device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
16model.to(device)
17model.eval()
18
19def generate_summary(input_text):
20 inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding="max_length", max_length=512).to(device)
21 with torch.no_grad():
22 summary_ids = model.generate(inputs["input_ids"], max_length=128, num_beams=4, early_stopping=True)
23 output_text = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
24 return output_text
25
26user_input = input("Enter your text: ")
27output = generate_summary(user_input)
28print("\nModel Output:")
29print(output)
30Training the Model
31The training process involves loading the pre-trained BART model and tokenizer, preparing a custom dataset, and training the model using the PyTorch DataLoader. Refer to the train_model() and evaluate_model() functions in the code for the detailed implementation.
32
33
34Feel free to modify any section to better fit your project’s needs!