ConvoBrief / ConvoBrief.py
Ketan Kumar
Add files via upload
dfda8e6 unverified
import transformers
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM,pipeline
from peft import PeftModel, PeftConfig
import streamlit as st
@st.cache_resource
def load_model():
config = PeftConfig.from_pretrained("Ketan3101/ConvoBrief")
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn")
model = PeftModel.from_pretrained(model, "Ketan3101/ConvoBrief")
tokenizer=AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
return model, tokenizer
def main():
st.set_page_config(page_title="ConvoBrief", page_icon="πŸ“")
model,tokenizer=load_model()
st.title("ConvoBrief: A dialogue summarizer")
dialogue=st.text_area("Enter the Dialogue")
if st.button("Summarize Dialogue"):
if dialogue:
inputs=tokenizer(dialogue,return_tensors='pt')
summary=tokenizer.decode(
model.generate(input_ids=inputs['input_ids'], max_new_tokens=200, temperature=1.2001, do_sample=True)[0],
skip_special_tokens=True
)
st.subheader("Summarized Dialogue:")
st.write(summary)
st.error("The model has been trained on less parameters, so their might be minor errors")
else:
st.warning("No! Dialogue was given")
if __name__=="__main__":
main()