Spaces:
Sleeping
Sleeping
| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| from conversation import get_default_conv_template | |
| import base64 | |
| import streamlit as st | |
| def generate(style, topic, words=200, sender='Sender_Name', recipient='Recipient_Name'): | |
| tokenizer = AutoTokenizer.from_pretrained("GeneZC/MiniChat-3B", use_fast=False) | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| def get_model(device): | |
| if device == 'cuda': | |
| return AutoModelForCausalLM.from_pretrained("GeneZC/MiniChat-3B", use_cache=True, device_map="auto", torch_dtype=torch.float16).eval() | |
| else: | |
| return AutoModelForCausalLM.from_pretrained("GeneZC/MiniChat-3B", use_cache=True, device_map="cpu", torch_dtype=torch.float32).eval() | |
| model = get_model(device) | |
| conv = get_default_conv_template("minichat") | |
| question = f""" | |
| Generate an email with the following specifications and make sure to highlight the subject, according to the topic of the mail, in the very first line of the response: | |
| - Style: {style} | |
| - Word Limit: {words} | |
| - Topic: {topic} | |
| - Sender Name: {sender} | |
| - Recipient Name: {recipient} | |
| """ | |
| conv.append_message(conv.roles[0], question) | |
| conv.append_message(conv.roles[1], None) | |
| prompt = conv.get_prompt() | |
| input_ids = tokenizer([prompt]).input_ids | |
| output_ids = model.generate( | |
| torch.as_tensor(input_ids).to(device), | |
| do_sample=True, | |
| temperature=0.7, | |
| max_new_tokens=2000, | |
| ) | |
| output_ids = output_ids[0][len(input_ids[0]):] | |
| output = tokenizer.decode(output_ids, skip_special_tokens=True).strip() | |
| return output | |
| def get_image(file): | |
| with open(file, "rb") as f: | |
| data = f.read() | |
| return base64.b64encode(data).decode() | |
| with open("./static/style.css") as style: | |
| st.markdown(f"<style>{style.read()}</style>", unsafe_allow_html=True) | |
| img = get_image("./static/bg.png") | |
| bg_image = f""" | |
| <style> | |
| [data-testid="stAppViewContainer"]{{ | |
| background-image: url("data:image/png;base64,{img}"); | |
| background-size: cover; | |
| }} | |
| </style> | |
| """ | |
| st.markdown(bg_image, unsafe_allow_html=True) | |
| st.title("MailCraft") | |
| st.subheader("Unlock Seamless Email Excellence 📧 for Effortless Communication") | |
| style = st.text_input("Enter Email Type", placeholder="Ex. Professional/Personal/Job Application") | |
| words = st.number_input("Enter Word Limit", min_value=100, max_value=500, value=None, placeholder="From 100 to 500") | |
| topic = st.text_input("Enter Email Topic") | |
| sender = st.text_input("Enter Sender Name") | |
| recipient = st.text_input("Enter Recipient Name") | |
| if st.button("Generate Email"): | |
| generated_email = generate(style, topic, words, sender, recipient) | |
| st.write(generated_email) |