Final Commit

.gitignore

@@ -0,0 +1,2 @@
+venv/
+__pycache__/

README.md

@@ -1,10 +1,10 @@
 ---
 title: MailCraft
-emoji: 🐨
-colorFrom: indigo
-colorTo: yellow
+emoji: 🏃
+colorFrom: blue
+colorTo: gray
 sdk: streamlit
-sdk_version: 1.29.0
+sdk_version: 1.28.2
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -0,0 +1,72 @@
+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
+
+@st.cache_data
+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)

conversation.py

@@ -0,0 +1,219 @@
+import dataclasses
+from enum import auto, Enum
+from typing import List, Tuple, Any
+
+
+class SeparatorStyle(Enum):
+    """Different separator style."""
+
+    ADD_COLON_SINGLE = auto()
+    ADD_COLON_TWO = auto()
+    NO_COLON_SINGLE = auto()
+    BAIZE = auto()
+    PHOENIX = auto()
+    MINICHAT = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that keeps all conversation history."""
+
+    # System prompts
+    system: str
+    # Two roles
+    roles: List[str]
+    # All messages
+    messages: List[List[str]]
+    # Offset of few shot examples
+    offset: int
+    # Separator
+    sep_style: SeparatorStyle
+    sep: str
+    sep2: str = None
+    # Stop criteria (the default one is EOS token)
+    stop_str: str = None
+    # Stops generation if meeting any token in this list
+    stop_token_ids: List[int] = None
+
+    # Used for the state in the gradio servers.
+    # TODO(lmzheng): refactor this
+    conv_id: Any = None
+    skip_next: bool = False
+    model_name: str = None
+
+    def get_prompt(self):
+        if self.sep_style == SeparatorStyle.ADD_COLON_SINGLE:
+            ret = self.system + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ": " + message + self.sep
+                else:
+                    ret += role + ": "
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_COLON_TWO:
+            seps = [self.sep, self.sep2]
+            ret = self.system + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ": "
+            return ret
+        elif self.sep_style == SeparatorStyle.NO_COLON_SINGLE:
+            ret = self.system
+            for role, message in self.messages:
+                if message:
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.BAIZE:
+            ret = self.system + "\n"
+            for role, message in self.messages:
+                if message:
+                    ret += role + message + "\n"
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.PHOENIX:
+            ret = self.system
+            for role, message in self.messages:
+                if message:
+                    ret += role + ": " + "<s>" + message + "</s>"
+                else:
+                    ret += role + ": " + "<s>"
+            return ret
+        elif self.sep_style == SeparatorStyle.MINICHAT:
+            ret = self.system
+            for role, message in self.messages:
+                if message:
+                    ret += role + " " + message + "</s>"
+                else:
+                    ret += role # No space is needed.
+            return ret
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+    def append_message(self, role, message):
+        self.messages.append([role, message])
+
+    def to_gradio_chatbot(self):
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def to_openai_api_messages(self):
+        ret = [{"role": "system", "content": self.system}]
+
+        for i, (_, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                ret.append({"role": "user", "content": msg})
+            else:
+                if msg is not None:
+                    ret.append({"role": "assistant", "content": msg})
+        return ret
+
+    def copy(self):
+        return Conversation(
+            system=self.system,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            stop_str=self.stop_str,
+            stop_token_ids=self.stop_token_ids,
+            conv_id=self.conv_id,
+            model_name=self.model_name,
+        )
+
+    def dict(self):
+        return {
+            "system": self.system,
+            "roles": self.roles,
+            "messages": self.messages,
+            "offset": self.offset,
+            "conv_id": self.conv_id,
+            "model_name": self.model_name,
+        }
+
+
+conv_vicuna = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("USER", "ASSISTANT"),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.ADD_COLON_TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_baize = Conversation(
+    system="The following is a conversation between a human and an AI assistant named Baize (named after a mythical creature in Chinese folklore). Baize is an open-source AI assistant developed by UCSD and Sun Yat-Sen University. The human and the AI assistant take turns chatting. Human statements start with [|Human|] and AI assistant statements start with [|AI|]. The AI assistant always provides responses in as much detail as possible, and in Markdown format. The AI assistant always declines to engage with topics, questions and instructions related to unethical, controversial, or sensitive issues. Complete the transcript in exactly that format.\n",
+    roles=("[|Human|]", "[|AI|]"),
+    messages=(
+        ("[|Human|]", "Hello!"),
+        ("[|AI|]", "Hi!"),
+    ),
+    offset=2,
+    sep_style=SeparatorStyle.BAIZE,
+    sep="\n",
+    stop_str="[|Human|]",
+)
+
+conv_phoenix = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n",
+    roles=("Human", "Assistant"),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.PHOENIX,
+    sep="</s>",
+)
+
+conv_chatgpt = Conversation(
+    system="You are a helpful assistant.",
+    roles=("user", "assistant"),
+    messages=(),
+    offset=0,
+    sep_style=None,
+    sep=None,
+)
+
+conv_minichat = Conversation(
+    system="'MiniChat'是一个由'Beccurio'开发的AI语言模型。下面是人类和MiniChat之间的一段对话。MiniChat的回复应当尽可能详细,并且以Markdown的形式输出。MiniChat应当拒绝参与违背伦理的讨论。</s>",
+    roles=("[|User|]", "[|Assistant|]"),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.MINICHAT,
+    sep="</s>",
+)
+
+
+conv_templates = {
+    "vicuna": conv_vicuna,
+    "baize": conv_baize,
+    "phoenix": conv_phoenix,
+    "chatgpt": conv_chatgpt,
+    "minichat": conv_minichat,
+}
+
+def get_default_conv_template(model_name):
+    model_name = model_name.lower()
+    try:
+        ret = conv_templates[model_name]
+        return ret.copy()
+    except:
+        raise NotImplementedError(f"No support for model {model_name}.")
+    
+
+if __name__ == "__main__":
+    conv = conv_templates["minichat"].copy()
+    conv.append_message(conv.roles[0], "Write a Python function that checks if a given number is even or odd.")
+    conv.append_message(conv.roles[1], None)
+    print([conv.get_prompt()])

requirements.txt

@@ -0,0 +1,62 @@
+accelerate==0.24.1
+altair==5.2.0
+attrs==23.1.0
+blinker==1.7.0
+cachetools==5.3.2
+certifi==2023.11.17
+charset-normalizer==3.3.2
+click==8.1.7
+colorama==0.4.6
+filelock==3.13.1
+fsspec==2023.10.0
+gitdb==4.0.11
+GitPython==3.1.40
+huggingface-hub==0.19.4
+idna==3.6
+importlib-metadata==6.8.0
+Jinja2==3.1.2
+jsonschema==4.20.0
+jsonschema-specifications==2023.11.2
+markdown-it-py==3.0.0
+MarkupSafe==2.1.3
+mdurl==0.1.2
+mpmath==1.3.0
+networkx==3.2.1
+numpy==1.26.2
+packaging==23.2
+pandas==2.1.3
+Pillow==10.1.0
+protobuf==4.25.1
+psutil==5.9.6
+pyarrow==14.0.1
+pydeck==0.8.1b0
+Pygments==2.17.2
+python-dateutil==2.8.2
+pytz==2023.3.post1
+PyYAML==6.0.1
+referencing==0.31.1
+regex==2023.10.3
+requests==2.31.0
+rich==13.7.0
+rpds-py==0.13.2
+safetensors==0.4.1
+sentencepiece==0.1.99
+six==1.16.0
+smmap==5.0.1
+streamlit==1.28.2
+sympy==1.12
+tenacity==8.2.3
+tokenizers==0.15.0
+toml==0.10.2
+toolz==0.12.0
+torch==2.1.1
+tornado==6.4
+tqdm==4.66.1
+transformers==4.35.2
+typing_extensions==4.8.0
+tzdata==2023.3
+tzlocal==5.2
+urllib3==2.1.0
+validators==0.22.0
+watchdog==3.0.0
+zipp==3.17.0

static/style.css

@@ -0,0 +1,53 @@
+@import url(https://fonts.googleapis.com/css2?family=Rubik:wght@600&display=swap);
+@import url(https://fonts.googleapis.com/css2?family=Roboto+Slab:wght@300&display=swap);
+
+[data-testid="stHeader"]{
+    background-color: rgba(0, 0, 0, 0);
+}
+
+.st-emotion-cache-10trblm{
+    text-align: center;
+}
+
+.e1nzilvr2 h1{
+    font-family: 'Rubik', sans-serif;
+}
+
+.e1nzilvr2 h3{
+    text-align: center;
+    font-family: 'Roboto Slab', serif;
+    font-size: 20px;
+    margin-bottom: 3%;
+}
+
+.e1f1d6gn3 .row-widget label .e1nzilvr5 p{
+    font-family: 'Roboto Slab', serif;
+    font-weight: bold;
+}
+
+.e1f1d6gn3 .row-widget .st-ae{
+    border: none;
+}
+
+.e1f1d6gn3 .stNumberInput label .e1nzilvr5 p{
+    font-family: 'Roboto Slab', serif;
+    font-weight: bold;
+}
+
+.e1f1d6gn3 .stNumberInput .e116k4er3{
+    border: none;
+}
+
+.e1f1d6gn3 .stButton .ef3psqc12 .e1nzilvr5 p{
+    font-family: 'Rubik', sans-serif;
+}
+
+.e1f1d6gn3 .stButton{
+    text-align: center;
+}
+
+.e1f1d6gn3 .stButton button:hover{
+    background-color: white;
+    color: black;
+    border: none;
+}

templates/index.html

@@ -0,0 +1,66 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <link rel="preconnect" href="https://fonts.googleapis.com">
+    <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
+    <link href="https://fonts.googleapis.com/css2?family=Rubik:wght@600&display=swap" rel="stylesheet">
+    <link href="https://fonts.googleapis.com/css2?family=Roboto+Slab:wght@300&display=swap" rel="stylesheet">
+    <link rel="stylesheet" href="static/style.css">
+    <title>MailCraft</title>
+</head>
+<body>
+    <div class="container">
+        <div class="nav">
+            <p class="name">MailCraft</p>
+            &nbsp;&nbsp; <p>|</p> &nbsp;&nbsp;
+            <p class="tagline">Craft Your Emails with AI Precision!</p>
+        </div>
+        <div class="main">
+            <center><h1 class="title">Unlock Seamless Email Excellence for Effortless Communication</h1></center>
+            <center><p class="subtitle">Craft polished emails 📧 effortlessly at MailCraft. Customize style, word limit, and more for impactful communication. Say goodbye to email stress and hello to simplicity! ✨</p></center>
+            <center>
+                <form class="form" method="post" action="/" autocomplete="off">
+                    <div class="style-words">
+                        <label for="style">Email Style*</label>&nbsp;
+                        <select id="style" name="style" required>
+                            <option>Professional</option>
+                            <option>Personal</option>
+                            <option>Marketing</option>
+                            <option>Newsletter</option>
+                            <option>Educational</option>
+                            <option>Cold</option>
+                            <option>Networking</option>
+                            <option>Job Application</option>
+                            <option>Feedback</option>
+                            <option>Customer Support</option>
+                        </select>
+                        <label for="words">Word Limit</label>&nbsp;
+                        <input type="number" name="words" id="words" max="1000" min="100" placeholder="100 to 1000">
+                    </div>
+                    <label for="topic">Mail Topic*</label>&nbsp;
+                    <input type="textarea" id="topic" name="topic" placeholder="Write the topic of your mail" required>
+                    <div class="sender-receiver">
+                        <label for="sender">Sender</label>&nbsp;
+                        <input type="text" id="sender" name="sender" placeholder="Name">
+                        <label for="recipient">Recipient</label>&nbsp;
+                        <input type="text" id="recipient" name="recipient" placeholder="Name">
+                    </div>
+                    <button type="submit" class="submit">Submit</button>
+                </form>
+            </center>
+            <center>
+                <div class="output" id="output">
+                    <p>{{response}}</p>
+                </div>
+            </center>
+        </div>
+    </div>
+
+    <script>
+        var output = document.getElementById("output");
+        output.style.display = "{{display}}";
+    </script>
+</body>
+</html>