Update app.py

app.py

@@ -1,8 +1,11 @@
 import gradio as gr
 import boto3
+import sagemaker
 import json
 import io
 import os
+from transformers import AutoTokenizer
+ 
 
 region = os.getenv("region")
 sm_endpoint_name = os.getenv("sm_endpoint_name")
@@ -16,20 +19,28 @@ session = boto3.Session(
 )
 sess = sagemaker.Session(boto_session=session)
 
+smr = session.client("sagemaker-runtime")
+
+DEFAULT_SYSTEM_PROMPT = (
+    "You are an helpful, concise and direct Assistant."
+)
+
+# load the tokenizer
+tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
+
+
+MAX_INPUT_TOKEN_LENGTH = 256
+
 # hyperparameters for llm
 parameters = {
     "do_sample": True,
     "top_p": 0.6,
     "temperature": 0.9,
-    "max_new_tokens": 1024,
+    "max_new_tokens": 768,
+    "repetition_penalty": 1.2,
     "return_full_text": False,
-    "stop": ["</s>"],
 }
 
-system_prompt = (
-    "You are an helpful Assistant, called Llama 2. Knowing everyting about AWS."
-)
-
 
 # Helper for reading lines from a stream
 class LineIterator:
@@ -61,66 +72,86 @@ class LineIterator:
             self.buffer.write(chunk["PayloadPart"]["Bytes"])
 
 
-# helper method to format prompt
-def create_messages_dict(message, history, system_prompt):
-    messages = []
-    if system_prompt:
-        messages.append({"role": "system", "content": system_prompt})
-    for user_prompt, bot_response in history:
-        messages.append({"role": "user", "content": user_prompt})
-        messages.append({"role": "assistant", "content": bot_response})
+def format_prompt(message, history):
+    '''
+    messages = [{"role": "system", "content": DEFAULT_SYSTEM_PROMPT}]
+    for interaction in history:
+        messages.append({"role": "user", "content": interaction[0]})
+        messages.append({"role": "assistant", "content": interaction[1]})
     messages.append({"role": "user", "content": message})
-    return messages
+    prompt = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    '''
 
+    messages = [
+    {"role": "user", "content": "Can you tell me an interesting fact about AWS?"},]
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+ 
+    return prompt
 
-def create_gradio_app(
-    endpoint_name,
-    session=boto3,
-    parameters=parameters,
-    system_prompt=system_prompt,
-    tokenizer=None,
-    concurrency_count=4,
-    share=True,
+
+def generate(
+    prompt,
+    history,
 ):
-    smr = session.client("sagemaker-runtime")
-
-    def generate(
-        prompt,
-        history,
-    ):
-        messages = create_messages_dict(prompt, history, system_prompt)
-        formatted_prompt = tokenizer.apply_chat_template(
-            messages, tokenize=False, add_generation_prompt=True
-        )
+    formatted_prompt = format_prompt(prompt, history)
+    check_input_token_length(formatted_prompt)
 
-        request = {"inputs": formatted_prompt, "parameters": parameters, "stream": True}
-        resp = smr.invoke_endpoint_with_response_stream(
-            EndpointName=endpoint_name,
-            Body=json.dumps(request),
-            ContentType="application/json",
-        )
+    request = {"inputs": formatted_prompt, "parameters": parameters, "stream": True}
 
-        output = ""
-        for c in LineIterator(resp["Body"]):
-            c = c.decode("utf-8")
-            if c.startswith("data:"):
-                chunk = json.loads(c.lstrip("data:").rstrip("/n"))
-                if chunk["token"]["special"]:
-                    continue
-                if chunk["token"]["text"] in request["parameters"]["stop"]:
-                    break
-                output += chunk["token"]["text"]
-                for stop_str in request["parameters"]["stop"]:
-                    if output.endswith(stop_str):
-                        output = output[: -len(stop_str)]
-                        output = output.rstrip()
-                        yield output
-
-                yield output
-        return output
-
-    demo = gr.ChatInterface(
-        generate, title="Chat with Amazon SageMaker", chatbot=gr.Chatbot(layout="panel")
+
+    resp = smr.invoke_endpoint_with_response_stream(
+        EndpointName=endpoint_name,
+        Body=json.dumps(request),
+        ContentType="application/json",
     )
+    output = ""
+    for c in LineIterator(resp["Body"]):
+        c = c.decode("utf-8")
+        if c.startswith("data:"):
+            chunk = json.loads(c.lstrip("data:").rstrip("/n"))
+            if chunk["token"]["special"]:
+                continue
+            if chunk["token"]["text"] in request["parameters"]["stop"]:
+                break
+            output += chunk["token"]["text"]
+            for stop_str in request["parameters"]["stop"]:
+                if output.endswith(stop_str):
+                    output = output[: -len(stop_str)]
+                    output = output.rstrip()
+                    yield output
+
+            yield output
+    return output
+
+
+def check_input_token_length(prompt: str) -> None:
+    input_token_length = len(tokenizer(prompt)["input_ids"])
+    if input_token_length > MAX_INPUT_TOKEN_LENGTH:
+        raise gr.Error(
+            f"The accumulated input is too long ({input_token_length} > {MAX_INPUT_TOKEN_LENGTH}). Clear your chat history and try again."
+        )
+
+
+theme = gr.themes.Monochrome(
+    primary_hue="indigo",
+    secondary_hue="blue",
+    neutral_hue="slate",
+    radius_size=gr.themes.sizes.radius_sm,
+    font=[
+        gr.themes.GoogleFont("Open Sans"),
+        "ui-sans-serif",
+        "system-ui",
+        "sans-serif",
+    ],
+)
+
+
+demo = gr.ChatInterface(
+    generate,
+    chatbot=gr.Chatbot(layout="panel"),
+    theme=theme,
+)
 
-    demo.queue(concurrency_count=concurrency_count).launch(share=share)
+demo.queue(concurrency_count=5).launch(share=False)