mail_fixer / app.py
tomaszki's picture
Add placeholder
56a4c1e
import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
class StreamlitStreamer(TextStreamer):
def on_finalized_text(self, text: str, stream_end: bool = False):
st.session_state['new_mail'] += text
new_mail.write(st.session_state['new_mail'])
device = 'cuda' if torch.cuda.is_available() else 'cpu'
@st.cache_resource
def load_model():
return AutoModelForCausalLM.from_pretrained(
"tomaszki/mail_fixer",
).to(device)
@st.cache_resource
def load_tokenizer():
return AutoTokenizer.from_pretrained("facebook/opt-125m")
model = load_model()
tokenizer = load_tokenizer()
st.title('Mail fixer')
placeholder = '''Hi Emma,
How did you like our art project? I thought it was pretty cool too! But one thing I noticed was that sometimes the letters were hard to see because they were a little bit too small. Let's try making them bigger or changing their shape so everyone can read them better.
Thanks for helping us work on the project! I have lots of other ideas for things we can do together. Can't wait to get started!
Love, James'''
mail = st.text_area('Enter your mail here', placeholder=placeholder, height=300)
new_mail = st.text('')
if mail:
st.session_state['new_mail'] = ''
streamer = StreamlitStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
prompt = f'Original email:\n{mail}\nFixed email:\n'
tokenized = tokenizer(prompt, return_tensors='pt').to(device)
output = model.generate(**tokenized, max_new_tokens=1024, streamer=streamer)