First app version

README.md

@@ -1,6 +1,6 @@
 ---
 title: Mail Fixer
-emoji: 🏃
+emoji: 📧
 colorFrom: yellow
 colorTo: blue
 sdk: streamlit

app.py

@@ -0,0 +1,36 @@
+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')
+
+mail = st.text_area('Enter your mail here')
+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)

requirements.txt

@@ -0,0 +1,4 @@
+transformers
+
+--find-links https://download.pytorch.org/whl/torch_stable.html
+torch==2.1.0+cpu