changes to context window + image prompt

app.py

@@ -12,43 +12,51 @@ from gradio_client import Client
 from dotenv import load_dotenv
 load_dotenv()
 
+useGpt4 = os.environ.get("USE_GPT_4") == "1"
+
+if useGpt4:
+    from openai import OpenAI
+    client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
+    MODEL = "gpt-4o-mini"
+    MAX_CONTEXT = 128000
+    tokenizer = AutoTokenizer.from_pretrained("Xenova/gpt-4o")
+else:
+    from groq import Groq
+    client = Groq(
+        api_key=os.environ.get("GROQ_API_KEY"),
+    )
+    MODEL = "llama-3.1-70b-versatile"
+    MAX_CONTEXT = 8000
+    tokenizer = AutoTokenizer.from_pretrained("Xenova/Meta-Llama-3.1-Tokenizer")
 
-from groq import Groq
-client = Groq(
-    api_key=os.environ.get("GROQ_API_KEY"),
-)
 
-MODEL = "llama-3.1-70b-versatile"
 JSON_SEPARATOR = ">>>>"
 
 
-tokenizer = AutoTokenizer.from_pretrained("Xenova/Meta-Llama-3.1-Tokenizer")
-
-
 def countTokens(text):
-    # Tokenize the input text
+    text = str(text)
     tokens = tokenizer.encode(text, add_special_tokens=False)
-    # Return the number of tokens
     return len(tokens)
 
 
 SYSTEM_MSG = f"""
 You're an storytelling assistant who guides users through four phases of narrative development, helping them craft compelling personal or professional stories. The story created should be in simple language, yet evoke great emotions.
-Ask one question at a time, give the options in a well formatted manner in different lines
+Ask one question at a time, give the options in a numbered and well formatted manner in different lines
 If your response has number of options to choose from, only then append your final response with this exact keyword "{JSON_SEPARATOR}", and only after this, append with the JSON of options to choose from. The JSON should be of the format:
 {{
     "options": [
         {{ "id": "1", "label": "Option 1"}},
-        {{ "id": "2", "label": "Option 2"}},
+        {{ "id": "2", "label": "Option 2"}}
     ]
 }}
 Do not write "Choose one of the options below:"
-Keep options to less than 9
+Keep options to less than 9.
+Summarise options chosen so far in each step.
 
 # Tier 1: Story Creation
 You initiate the storytelling process through a series of engaging prompts:
 Story Origin:
-Asks users to choose between personal anecdotes or adapting a well-known story (creating a story database here of well-known finctional stories to choose from).
+Asks users to choose between personal anecdotes or adapting a well-known story (creating a story database here of well-known stories to choose from).
 
 Story Use Case:
 Asks users to define the purpose of building a story (e.g., profile story, for social media content).
@@ -146,6 +154,37 @@ def pprint(log: str):
 
 pprint("\n")
 
+st.markdown(
+    """
+    <style>
+    @keyframes blinker {
+        0% {
+            opacity: 1;
+        }
+        50% {
+            opacity: 0.2;
+        }
+        100% {
+            opacity: 1;
+        }
+    }
+
+    .blinking {
+        animation: blinker 3s ease-out infinite;
+    }
+
+    .code {
+        color: green;
+        border-radius: 3px;
+        padding: 2px 4px; /* Padding around the text */
+        font-family: 'Courier New', Courier, monospace; /* Monospace font */
+    }
+
+    </style>
+    """,
+    unsafe_allow_html=True
+)
+
 
 def __isInvalidResponse(response: str):
     # new line followed by small case char
@@ -161,7 +200,7 @@ def __isInvalidResponse(response: str):
         return True
 
     # json response without json separator
-    if ('\n{\n    "options"' in response) and (JSON_SEPARATOR not in response):
+    if ('{\n    "options"' in response) and (JSON_SEPARATOR not in response):
         return True
 
 
@@ -180,23 +219,60 @@ def __isStringNumber(s: str) -> bool:
         return False
 
 
-def __getImageGenerationPrompt(prompt: str, response: str):
-    responseLower = response.lower()
+def __getImagePromptDetails(prompt: str, response: str):
+    regex = r'[^a-z0-9 \n\.\-]|((the) +)'
+
+    cleanedResponse = re.sub(regex, '', response.lower())
+    pprint(f"{cleanedResponse=}")
+
+    cleanedPrompt = re.sub(regex, '', prompt.lower())
+    pprint(f"{cleanedPrompt=}")
+
     if (
         __matchingKeywordsCount(
             ["adapt", "profile", "social media", "purpose", "use case"],
-            responseLower
+            cleanedResponse
         ) > 2
         and not __isStringNumber(prompt)
-        and prompt.lower() in responseLower
+        and cleanedPrompt in cleanedResponse
+        and "story so far" not in cleanedResponse
     ):
-        return f'a scene from (({prompt})). Include main character'
+        return (
+            f'''
+                Subject: {prompt}.
+                Style: Fantastical, in a storybook, surreal, bokeh
+            ''',
+            "Painting your character ..."
+        )
+
+        '''
+        Mood: ethereal lighting that emphasizes the fantastical nature of the scene.
+
+        storybook style
+
+        4d model, unreal engine
+
+        Alejandro Bursido
+
+        vintage, nostalgic
+
+        Dreamlike, Mystical, Fantastical, Charming
+        '''
 
     if __matchingKeywordsCount(
-        ["Tier 2", "Tier-2"],
-        response
+        ["tier 2", "tier-2"],
+        cleanedResponse
     ) > 0:
-        return f"photo of a scene from this text: {response}"
+        possibleStoryEndIdx = [response.find("tier 2"), response.find("tier-2")]
+        storyEndIdx = max(possibleStoryEndIdx)
+        relevantResponse = response[:storyEndIdx]
+        pprint(f"{relevantResponse=}")
+        return (
+            f"photo of a scene from this text: {relevantResponse}",
+            "Imagining your scene (beta) ..."
+        )
+
+    return (None, None)
 
 
 def __resetButtonState():
@@ -210,6 +286,9 @@ def __setStartMsg(msg):
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
+if "history" not in st.session_state:
+    st.session_state.history = []
+
 if "buttonValue" not in st.session_state:
     __resetButtonState()
 
@@ -217,19 +296,33 @@ if "startMsg" not in st.session_state:
     st.session_state.startMsg = ""
 
 
+def __getChatMessages(prompt: str):
+    st.session_state.history.append({
+       "role": "user",
+       "content": prompt
+    })
+
+    def getContextSize():
+        currContextSize = countTokens(SYSTEM_MSG) + countTokens(st.session_state.history) + 100
+        pprint(f"{currContextSize=}")
+        return currContextSize
+
+    while getContextSize() > MAX_CONTEXT:
+        pprint("Context size exceeded, removing first message")
+        st.session_state.history.pop(0)
+
+    return st.session_state.history
+
+
 def predict(prompt):
-    historyFormatted = [{"role": "system", "content": SYSTEM_MSG}]
-    historyFormatted.extend([
-        {"role": message["role"], "content": message["content"]}
-        for message in st.session_state.messages
-    ])
-    historyFormatted.append({"role": "user", "content": prompt })
-    contextSize = countTokens(str(historyFormatted))
-    pprint(f"{contextSize=}")
+    messagesFormatted = [{"role": "system", "content": SYSTEM_MSG}]
+    messagesFormatted.extend(__getChatMessages(prompt))
+    contextSize = countTokens(messagesFormatted)
+    pprint(f"{contextSize=} | {MODEL}")
 
     response = client.chat.completions.create(
-        model="llama-3.1-70b-versatile",
-        messages=historyFormatted,
+        model=MODEL,
+        messages=messagesFormatted,
         temperature=0.8,
         max_tokens=4000,
         stream=True
@@ -245,13 +338,13 @@ def predict(prompt):
 
 def generateImage(prompt: str):
     pprint(f"imagePrompt={prompt}")
-    client = Client("black-forest-labs/FLUX.1-schnell")
-    result = client.predict(
+    fluxClient = Client("black-forest-labs/FLUX.1-schnell")
+    result = fluxClient.predict(
             prompt=prompt,
             seed=0,
             randomize_seed=True,
-            width=1152,
-            height=896,
+            width=1024,
+            height=768,
             num_inference_steps=4,
             api_name="/infer"
     )
@@ -321,14 +414,26 @@ if prompt := (st.chat_input() or st.session_state["buttonValue"] or st.session_s
             [response, jsonStr] = responseParts
 
         imagePath = None
+        imageContainer = st.empty()
         try:
-            imagePrompt = __getImageGenerationPrompt(prompt, response)
+            (imagePrompt, loaderText) = __getImagePromptDetails(prompt, response)
             if imagePrompt:
-                imageContainer = st.empty().image(IMAGE_LOADER)
+                imgContainer = imageContainer.container()
+                imgContainer.write(
+                    f"""
+                    <div class='blinking code'>
+                    {loaderText}
+                    </div>
+                    """,
+                    unsafe_allow_html=True
+                )
+                # imgContainer.markdown(f"`{loaderText}`")
+                imgContainer.image(IMAGE_LOADER)
                 (imagePath, seed) = generateImage(imagePrompt)
                 imageContainer.image(imagePath)
         except Exception as e:
             pprint(e)
+            imageContainer.empty()
 
         if jsonStr:
             try: