Update app.py

Adding default text, instruct prefix, and more examples

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-import os, gc
+import os, gc, copy
 from huggingface_hub import hf_hub_download
 from pynvml import *
 
@@ -39,15 +39,70 @@ model = RWKV(model=model_path, strategy=MODEL_STRAT)
 from rwkv.utils import PIPELINE
 pipeline = PIPELINE(model, "rwkv_vocab_v20230424")
 
+# Precomputation of the state
+def precompute_state(text):
+    state = None
+    text_encoded = pipeline.encode(text)
+    _, state = model.forward(text_encoded)
+    yield state
+
+# Precomputing the base instruction set
+INSTRUCT_PREFIX = f'''
+The following is a set of instruction rules, that can translate spoken text to zombie speak. And vice visa.
+
+# Zombie Speak Rules:
+- Replace syllables with "uh" or "argh"
+- Add "uh" and "argh" sounds between words
+- Repeat words and letters, especially vowels
+- Use broken grammar and omit small words like "the", "a", "is"
+
+# To go from zombie speak back to English:
+- Remove extra "uh" and "argh" sounds
+- Replace repeated letters with one instance
+- Add omitted small words like "the", "a", "is" back in
+- Fix grammar and sentence structure
+
+# Here are several examples:
+
+## English: 
+"Hello my friend, how are you today?"
+## Zombie: 
+"Hell-uh-argh myuh fruh-end, hargh-owuh argh yuh-uh toduh-ay?"
+
+## Zombie: 
+"Brargh-ains argh-uh foo-duh"
+## English: 
+"Brains are food"
+
+## English:
+"Good morning! How are you today? I hope you are having a nice day. The weather is supposed to be sunny and warm this afternoon. Maybe we could go for a nice walk together and stop to get ice cream. That would be very enjoyable. Well, I will talk to you soon!"
+## Zombie:
+"Guh-ood morngh-ing! Hargh-owuh argh yuh-uh toduh-ay? Iuh hargh-ope yuh-uh argh havi-uh-nguh nuh-ice duh-ay. Thuh-uh weath-uh-eruh izzuh suh-pose-duh tuh-uh beh sunn-eh an-duh war-muh thizuh aft-erng-oon. May-buh-uh weh coulduh gargh-oh fargh-oruh nuh-ice wal-guh-kuh toge-the-ruh an-duh stargh-op tuh-uh geh-etuh izz-creem. Tha-at wou-duh beh ve-reh uhn-joy-ab-buhl. Well, I wih-ll targh-alk tuh-uh yuh-oo soo-oon!"
+
+'''
+
+# Get the prefix state
+PREFIX_STATE = precompute_state(INSTRUCT_PREFIX)
+
 # Translation logic
-def translate(text, target_language):
+def translate(text, target_language, inState=PREFIX_STATE):
     prompt = f"Translate the following text to {target_language}\n # Input Text:\n{text}\n\n# Output Text:\n"
     ctx = prompt.strip()
     all_tokens = []
     out_last = 0
     out_str = ''
     occurrence = {}
+
     state = None
+    if inState != None:
+        state = copy.deepcopy(inState)
+    
+    # Clear GC
+    gc.collect()
+    if HAS_GPU == True :
+        torch.cuda.empty_cache()
+
+    # Generate things token by token
     for i in range(ctx_limit):
         out, state = model.forward(pipeline.encode(ctx)[-ctx_limit:] if i == 0 else [token], state)
         token = pipeline.sample_logits(out)
@@ -60,19 +115,24 @@ def translate(text, target_language):
             yield out_str.strip()
             out_last = i + 1
 
+        if "# " in out_str and "\n#" in out_str 
+            out_str = out_str.split("\n## ")[0].split("\n# ")[0]
+            yield out_str.strip()
+
     del out
     del state
 
-    # Clear GC
-    gc.collect()
-    if HAS_GPU == True :
-        torch.cuda.empty_cache()
+    # # Clear GC
+    # gc.collect()
+    # if HAS_GPU == True :
+    #     torch.cuda.empty_cache()
     
     yield out_str.strip()
 
 # Languages
 LANGUAGES = [
   "English",
+  "Zombie Speak",
   "Chinese",
   "Spanish",
   "Bengali",
@@ -183,6 +243,8 @@ LANGUAGES = [
 
 # Example data
 EXAMPLES = [
+    ["Brargh-ains argh-uh foo-duh", "English"],
+    ["I Want to eat your brains", "Zombie Speak"],
     ["Hello, how are you?", "French"],
     ["Hello, how are you?", "Spanish"],
     ["Hello, how are you?", "Chinese"],
@@ -197,11 +259,11 @@ EXAMPLES = [
 # Gradio interface
 with gr.Blocks(title=title) as demo:
     gr.HTML(f"<div style=\"text-align: center;\"><h1>RWKV-5 World v2 - {title}</h1></div>")
-    gr.Markdown("This is the RWKV-5 World v2 1B5 model tailored for translation. Please provide the text and select the target language for translation.")
+    gr.Markdown("This is the RWKV-5 World v2 1B5 model tailored for translation. With a halloween zombie speak twist")
     
     # Input and output components
-    text = gr.Textbox(lines=5, label="Source Text", placeholder="Enter the text you want to translate...")
-    target_language = gr.Dropdown(choices=LANGUAGES, label="Target Language")
+    text = gr.Textbox(lines=5, label="Source Text", placeholder="Enter the text you want to translate...", default=EXAMPLES[0][0])
+    target_language = gr.Dropdown(choices=LANGUAGES, label="Target Language", default=EXAMPLES[0][1])
     output = gr.Textbox(lines=5, label="Translated Text")
     submit = gr.Button("Translate", variant="primary")