Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from transformers import ( | |
| AutoModelForSeq2SeqLM, # Import the Seq2Seq language model class | |
| AutoTokenizer, # Import the tokenizer for language models | |
| GenerationConfig, # Import the generation configuration for text generation | |
| TrainingArguments, # Import training arguments for fine-tuning models | |
| Trainer # Import the Trainer class for model training | |
| ) | |
| import torch | |
| # Import the necessary components for PEFT model | |
| from peft import PeftModel, PeftConfig | |
| # Load the base T5 model for the PEFT model | |
| peft_model_base = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base", torch_dtype=torch.bfloat16) | |
| tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base") | |
| # Create the PEFT model by loading the pretrained base model and checkpoint | |
| peft_model = PeftModel.from_pretrained( | |
| peft_model_base, # Base model | |
| './peft-dialogue-summary-checkpoint-local/', # Checkpoint directory | |
| torch_dtype=torch.bfloat16, # PyTorch data type | |
| is_trainable=False # Specify that the model is not trainable | |
| ) | |
| st.title("Dialogue Summarization using Flan T5 model") | |
| with st.expander("Info about Algorithm and Dataset"): | |
| st.markdown(""" Open-source [FlanT5 model](https://huggingface.co/google/flan-t5-base) | |
| has been finetuned on Open-source [data](https://huggingface.co/datasets/knkarthick/dialogsum) | |
| using PEFT (Parameter efficient fine-tuning) for Dialogue Summarization """) | |
| dialogue = st.text_area( | |
| "Dialogue for Summarization", | |
| """ #Person1#: Have you considered upgrading your system? | |
| #Person2#: Yes, but I'm not sure what exactly I would need. | |
| #Person1#: You could consider adding a painting program to your software. It would allow you to make up your own flyers and banners for advertising. | |
| #Person2#: That would be a definite bonus. | |
| #Person1#: You might also want to upgrade your hardware because it is pretty outdated now. | |
| #Person2#: How can we do that? | |
| #Person1#: You'd probably need a faster processor, to begin with. And you also need a more powerful hard disc, more memory and a faster modem. Do you have a CD-ROM drive? | |
| #Person2#: No. | |
| #Person1#: Then you might want to add a CD-ROM drive too, because most new software programs are coming out on Cds. | |
| #Person2#: That sounds great. Thanks. """, height=323) | |
| prompt = f""" Summarize the following conversation. | |
| {dialogue} | |
| Summary: """ | |
| if st.button("Summarize", type="primary"): | |
| # Tokenize the prompt and convert it to PyTorch tensors | |
| input_ids = tokenizer(prompt, return_tensors="pt").input_ids | |
| peft_model_outputs = peft_model.generate(input_ids=input_ids, generation_config=GenerationConfig(max_new_tokens=200, num_beams=1)) | |
| peft_model_text_output = tokenizer.decode(peft_model_outputs[0], skip_special_tokens=True) | |
| st.text(peft_model_text_output) | |