Summarization / app.py
ketanchaudhary88's picture
Create app.py
3037b87 verified
raw
history blame
1.52 kB
from transformers import BartForConditionalGeneration, BartTokenizer
# Load the pre-trained BART model and tokenizer
model_name = "facebook/bart-large-cnn"
model = BartForConditionalGeneration.from_pretrained(model_name)
tokenizer = BartTokenizer.from_pretrained(model_name)
# Define a function to summarize a conversation
def summarize_conversation(conversation):
# Join the conversation into a single text (separate sentences by a space)
conversation_text = " ".join(conversation)
# Tokenize the conversation
inputs = tokenizer(conversation_text, return_tensors="pt", max_length=1024, truncation=True)
# Generate the summary
summary_ids = model.generate(inputs['input_ids'], num_beams=4, min_length=50, max_length=200, early_stopping=True)
# Decode the summary
summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
return summary
# Sample conversation between agent and customer
conversation = [
"Hello, how can I assist you today?",
"I need help with my order. I haven't received it yet.",
"Can you provide me with your order number?",
"Sure, my order number is 123456.",
"Thank you. Let me check the status of your order.",
"It looks like your order was delayed due to some shipping issues.",
"When will I receive my order?",
"The order is expected to arrive within the next 2-3 days."
]
# Generate summary
summary = summarize_conversation(conversation)
# Print the summary
print("Conversation Summary:")
print(summary)