Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-import os, gc, copy
+import os, gc, copy, torch
 from huggingface_hub import hf_hub_download
 from pynvml import *
 
@@ -39,45 +39,54 @@ model = RWKV(model=model_path, strategy=MODEL_STRAT)
 from rwkv.utils import PIPELINE
 pipeline = PIPELINE(model, "rwkv_vocab_v20230424")
 
+# State copy
+def universal_deepcopy(obj):
+    if hasattr(obj, 'clone'):  # Assuming it's a tensor if it has a clone method
+        return obj.clone()
+    elif isinstance(obj, list):
+        return [universal_deepcopy(item) for item in obj]
+    else:
+        return copy.deepcopy(obj)
+
+# For debgging mostly
+def inspect_structure(obj, depth=0):
+    indent = "  " * depth
+    obj_type = type(obj).__name__
+    
+    if isinstance(obj, list):
+        print(f"{indent}List (length {len(obj)}):")
+        for item in obj:
+            inspect_structure(item, depth + 1)
+    elif isinstance(obj, dict):
+        print(f"{indent}Dict (length {len(obj)}):")
+        for key, value in obj.items():
+            print(f"{indent}  Key: {key}")
+            inspect_structure(value, depth + 1)
+    else:
+        print(f"{indent}{obj_type}")
+
 # Precomputation of the state
 def precompute_state(text):
     state = None
     text_encoded = pipeline.encode(text)
     _, state = model.forward(text_encoded, state)
-    yield dict(state)
+    return 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"
+You are a translator bot that can translate text to any language.
+And will respond only with the translated text, without additional comments.
 
-# 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
+## From English:
+It is not enough to know, we must also apply; it is not enough to will, we must also do.
+## To Polish:
+Nie wystarczy wiedzieć, trzeba także zastosować; nie wystarczy chcieć, trzeba też działać.
 
-# Here are several examples:
+## From Spanish:
+La muerte no nos concierne, porque mientras existamos, la muerte no está aquí. Y cuando llega, ya no existimos.
+## To English:
+Death does not concern us, because as long as we exist, death is not here. And when it does come, we no longer exist.
 
-## 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!"
 
 '''
 
@@ -85,8 +94,8 @@ The following is a set of instruction rules, that can translate spoken text to z
 PREFIX_STATE = precompute_state(INSTRUCT_PREFIX)
 
 # Translation logic
-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"
+def translate(text, source_language, target_language, inState=PREFIX_STATE):
+    prompt = f"## From {source_language}:\n{text}\n\n## To {target_language}:\n"
     ctx = prompt.strip()
     all_tokens = []
     out_last = 0
@@ -95,7 +104,7 @@ def translate(text, target_language, inState=PREFIX_STATE):
 
     state = None
     if inState != None:
-        state = dict(inState)
+        state = universal_deepcopy(inState)
     
     # Clear GC
     gc.collect()
@@ -105,19 +114,43 @@ def translate(text, target_language, inState=PREFIX_STATE):
     # 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)
+        token = pipeline.sample_logits(out, temperature=0.1, top_p=0.5)
         if token in [0]:  # EOS token
             break
         all_tokens += [token]
         tmp = pipeline.decode(all_tokens[out_last:])
         if '\ufffd' not in tmp:
             out_str += tmp
-            yield out_str.strip()
             out_last = i + 1
+        else:
+            return out_str.strip()
+
+        if "\n:" in out_str :
+            out_str = out_str.split("\n\nHuman:")[0].split("\nHuman:")[0]
+            return out_str.strip()
 
-        if "# " in out_str and "\n#" in out_str :
-            out_str = out_str.split("\n## ")[0].split("\n# ")[0]
-            yield out_str.strip()
+        if "{source_language}:" in out_str :
+            out_str = out_str.split("{source_language}:")[0]
+            return out_str.strip()
+
+        if "{target_language}:" in out_str :
+            out_str = out_str.split("{target_language}:")[0]
+            return out_str.strip()
+
+        if "\nHuman:" in out_str :
+            out_str = out_str.split("\n\nHuman:")[0].split("\nHuman:")[0]
+            return out_str.strip()
+
+        if "\nAssistant:" in out_str :
+            out_str = out_str.split("\n\nAssistant:")[0].split("\nAssistant:")[0]
+            return out_str.strip()
+
+        if "\n#" in out_str :
+            out_str = out_str.split("\n\n#")[0].split("\n#")[0]
+            return out_str.strip()
+
+        # Yield for streaming
+        yield out_str.strip()
 
     del out
     del state
@@ -127,7 +160,8 @@ def translate(text, target_language, inState=PREFIX_STATE):
     # if HAS_GPU == True :
     #     torch.cuda.empty_cache()
     
-    yield out_str.strip()
+    # yield out_str.strip()
+    return out_str.strip()
 
 # Languages
 LANGUAGES = [
@@ -219,6 +253,7 @@ LANGUAGES = [
   "Arabic (Tunisian Spoken)",
   "Rwanda",
   "Zulu",
+  "Latin",
   "Bulgarian",
   "Swedish",
   "Lombard",
@@ -243,37 +278,34 @@ 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"],
-    ["Bonjour, comment ça va?", "English"],
-    ["Hola, ¿cómo estás?", "English"],
-    ["你好吗?", "English"],
-    ["Guten Tag, wie geht es Ihnen?", "English"],
-    ["Привет, как ты?", "English"],
-    ["مرحبًا ، كيف حالك؟", "English"],
+    # More people would learn from their mistakes if they weren't so busy denying them.
+    ["Többen tanulnának a hibáikból, ha nem lennének annyira elfoglalva, hogy tagadják azokat.", "Hungarian", "English"],
+    ["La mejor venganza es el éxito masivo.", "Spanish", "English"],
+    ["Tout est bien qui finit bien.", "French", "English"],
+    ["Lasciate ogne speranza, voi ch'intrate.", "Italian", "English"],
+    ["Errare humanum est.", "Latin", "English"],
 ]
-
 # 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. With a halloween zombie speak twist")
+    gr.Markdown("This is the RWKV-5 World v2 1B5 model tailored for translation tasks")
     
     # Input and output components
-    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])
+    text = gr.Textbox(lines=5, label="Source Text", placeholder="Enter the text you want to translate...", value=EXAMPLES[0][0])
+    source_language = gr.Dropdown(choices=LANGUAGES, label="Source Language", value=EXAMPLES[0][1])
+    target_language = gr.Dropdown(choices=LANGUAGES, label="Target Language", value=EXAMPLES[0][2])
     output = gr.Textbox(lines=5, label="Translated Text")
+
+    # Submission
     submit = gr.Button("Translate", variant="primary")
     
     # Example data
-    data = gr.Dataset(components=[text, target_language], samples=EXAMPLES, label="Example Translations", headers=["Text", "Target Language"])
+    data = gr.Dataset(components=[text, source_language, target_language], samples=EXAMPLES, label="Example Translations", headers=["Source Text", "Target Language"])
     
     # Button action
-    submit.click(translate, [text, target_language], [output])
-    data.click(lambda x: x, [data], [text, target_language])
+    submit.click(translate, [text, source_language, target_language], [output])
+    data.click(lambda x: x, [data], [text, source_language, target_language])
 
 # Gradio launch
 demo.queue(concurrency_count=1, max_size=10)
-demo.launch(share=False)
+demo.launch(share=False, debug=True)