修改句子輸出，看起來整齊點

README.md

@@ -1,10 +1,10 @@
 ---
 title: SD_Helper_01
 emoji: 📊
-colorFrom: blue
-colorTo: green
+colorFrom: gray
+colorTo: indigo
 sdk: gradio
-sdk_version: 3.24.1
+sdk_version: 3.30.0
 app_file: app.py
 pinned: false
 license: openrail

app.py

@@ -27,51 +27,44 @@ zh2en_tokenizer = AutoTokenizer.from_pretrained('Helsinki-NLP/opus-mt-zh-en')
 en2zh_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-zh").eval()
 en2zh_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-zh")
 
-def load_prompter():
-    prompter_model = AutoModelForCausalLM.from_pretrained("microsoft/Promptist")
-    tokenizer = AutoTokenizer.from_pretrained("gpt2")
-    tokenizer.pad_token = tokenizer.eos_token
-    tokenizer.padding_side = "left"
-    return prompter_model, tokenizer
 
-prompter_model, prompter_tokenizer = load_prompter()
+def translate_zh2en(text):
+    with torch.no_grad():
 
-def generate_prompter(plain_text, max_new_tokens=75, num_return_sequences=3):
-    input_ids = prompter_tokenizer(plain_text.strip() + " Rephrase:", return_tensors="pt").input_ids
-    eos_id = prompter_tokenizer.eos_token_id
-    outputs = prompter_model.generate(
-        input_ids,
-        do_sample=False,
-        max_new_tokens=75,
-        num_beams=6,
-        num_return_sequences=num_return_sequences,
-        eos_token_id=eos_id,
-        pad_token_id=eos_id,
-        length_penalty=-1
+        text = re.sub(r'([^\u4e00-\u9fa5])([\u4e00-\u9fa5])', r'\1\n\2', text)
+        text = re.sub(r'([\u4e00-\u9fa5])([^\u4e00-\u9fa5])', r'\1\n\2', text)
 
-    )
+        text = text.replace('\n', ',')
 
-    output_texts = prompter_tokenizer.batch_decode(outputs, skip_special_tokens=True)
-    result = ""
-    for output_text in output_texts:
-        result.append(output_text.replace(plain_text + " Rephrase:", "").strip())
+        text =re.sub(r'(?<![a-zA-Z])\s+|\s+(?![a-zA-Z])', '', text)
 
-    return "\n".join(result)
+        text = re.sub(r',+', ',', text)
 
-def translate_zh2en(text):
-    with torch.no_grad():
-        text = text.replace('\n', ',').replace('\r', ',')
-        text = re.sub('^,+', ',', text)
         encoded = zh2en_tokenizer([text], return_tensors='pt')
         sequences = zh2en_model.generate(**encoded)
-        return zh2en_tokenizer.batch_decode(sequences, skip_special_tokens=True)[0]
+        result = zh2en_tokenizer.batch_decode(sequences, skip_special_tokens=True)[0]
+
+        result = result.strip() 
+
+        return result
+
 
 def translate_en2zh(text):
     with torch.no_grad():
+
         encoded = en2zh_tokenizer([text], return_tensors="pt")
         sequences = en2zh_model.generate(**encoded)
         return en2zh_tokenizer.batch_decode(sequences, skip_special_tokens=True)[0]
 
+def test05(text):
+
+    return text
+
+def test06(text):
+
+    return text
+
+
 def text_generate(text):
     seed = random.randint(100, 1000000)
     set_seed(seed)
@@ -83,53 +76,118 @@ def text_generate(text):
         list = []
         for sequence in sequences:
 
+
             line = sequence['generated_text'].strip()
 
-            if line != text_in_english and len(line) > (len(text_in_english) + 4) and line.endswith(
-                    (':', '-', '—')) is False:
-                list.append(line)
+            if line != text_in_english and len(line) > (len(text_in_english) + 4):
+
+                list.append(translate_en2zh(line)+"\n")
+                list.append(line+"\n")
+                list.append("\n")
 
-        result = "\n".join(list)
+        result = "".join(list)
 
         result = re.sub('[^ ]+\.[^ ]+', '', result)
 
-        result = result.replace('<', '').replace('>', '').replace('"', '')
+        result = result.replace('<', '').replace('>', '')
+
         if result != '':
             break
+    return result
+
+
+def load_prompter():
+    prompter_model = AutoModelForCausalLM.from_pretrained("microsoft/Promptist")
+    tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.padding_side = "left"
+    return prompter_model, tokenizer
+
+prompter_model, prompter_tokenizer = load_prompter()
+
+def generate_prompter(text):
+    text = translate_zh2en(text)
+
+    input_ids = prompter_tokenizer(text.strip()+" Rephrase:", return_tensors="pt").input_ids
+    eos_id = prompter_tokenizer.eos_token_id
+    outputs = prompter_model.generate(
+        input_ids,
+        do_sample=False,
+        max_new_tokens=75,
+        num_beams=3,
+        num_return_sequences=3,
+        eos_token_id=eos_id,
+        pad_token_id=eos_id,
+        length_penalty=-1.0
+    )
+    output_texts = prompter_tokenizer.batch_decode(outputs, skip_special_tokens=True)
+
+    result = []
+    for output_text in output_texts:
 
-    return result, "\n".join(translate_en2zh(line) for line in result.split("\n") if len(line) > 0) 
+        output_text = output_text.replace('<', '').replace('>', '')
+        output_text = output_text.split("Rephrase:", 1)[-1].strip()
+
+        result.append(translate_en2zh(output_text)+"\n")
+        result.append(output_text+"\n")
+        result.append("\n")
+    return "".join(result)
+
+def combine_text(text):
+    text01 = generate_prompter(text)
+    text02 = text_generate(text)
+    return text01,text02
 
 def get_prompt_from_image(input_image):
     image = input_image.convert('RGB')
     pixel_values = big_processor(images=image, return_tensors="pt").to(device).pixel_values
     generated_ids = big_model.to(device).generate(pixel_values=pixel_values, max_length=50)
     generated_caption = big_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    print(generated_caption)
-    return generated_caption
-
+    result01 = generate_prompter(generated_caption)
+    result02 = text_generate(generated_caption)
+    return result01,result02
 
 with gr.Blocks() as block:
     with gr.Column():
-        with gr.Tab('文生文'):
+        with gr.Tab('工作區'):
             with gr.Row():
                 input_text = gr.Textbox(lines=12, label='輸入文字', placeholder='在此输入文字...')
-
+                input_image = gr.Image(type='pil')
             with gr.Row():
-                txt_prompter_btn = gr.Button('執行')
-
-        with gr.Tab('圖生文'):
+                txt_prompter_btn = gr.Button('文生文')
+                pic_prompter_btn = gr.Button('圖生文')
             with gr.Row():
-                input_image = gr.Image(type='pil')
-
+                Textbox_1 = gr.Textbox(lines=6, label='生成方式A')
             with gr.Row():
-                pic_prompter_btn = gr.Button('執行')
-
-    Textbox_1 = gr.Textbox(lines=6, label='輸出結果')
-    Textbox_2 = gr.Textbox(lines=6, label='中文翻譯')
+                Textbox_2 = gr.Textbox(lines=6, label='生成方式B')
+        with gr.Tab('測試區'):
+            with gr.Row():
+                input_test01 = gr.Textbox(lines=2, label='中英翻譯', placeholder='在此输入文字...')
+                test01_btn = gr.Button('執行')
+                Textbox_test01 = gr.Textbox(lines=2, label='輸出結果')
+            with gr.Row():
+                input_test02 = gr.Textbox(lines=2, label='英中翻譯', placeholder='在此输入文字...')
+                test02_btn = gr.Button('執行')
+                Textbox_test02 = gr.Textbox(lines=2, label='輸出結果')
+            with gr.Row():
+                input_test03 = gr.Textbox(lines=2, label='SD模式', placeholder='在此输入文字...')
+                test03_btn = gr.Button('執行')
+                Textbox_test03 = gr.Textbox(lines=2, label='輸出結果')
+            with gr.Row():
+                input_test04 = gr.Textbox(lines=2, label='瞎掰模式', placeholder='在此输入文字...')
+                test04_btn = gr.Button('執行')
+                Textbox_test04 = gr.Textbox(lines=2, label='輸出結果')
+            with gr.Row():
+                input_test05 = gr.Textbox(lines=2, label='沒作用', placeholder='在此输入文字...')
+                test05_btn = gr.Button('執行')
+                Textbox_test05 = gr.Textbox(lines=2, label='輸出結果')
+            with gr.Row():
+                input_test06 = gr.Textbox(lines=2, label='沒作用', placeholder='在此输入文字...')
+                test06_btn = gr.Button('執行')
+                Textbox_test06 = gr.Textbox(lines=2, label='輸出結果')
 
     txt_prompter_btn.click(
-
-        fn=text_generate,
+        fn=combine_text,
         inputs=input_text, 
         outputs=[Textbox_1,Textbox_2]
     )
@@ -137,7 +195,41 @@ with gr.Blocks() as block:
     pic_prompter_btn.click(
         fn=get_prompt_from_image,
         inputs=input_image,
-        outputs=Textbox_1
+        outputs=[Textbox_1,Textbox_2]
+    )
+
+    test01_btn.click(
+        fn=translate_zh2en,
+        inputs=input_test01,
+        outputs=Textbox_test01
+    )
+
+    test02_btn.click(
+        fn=translate_en2zh,
+        inputs=input_test02,
+        outputs=Textbox_test02
+    )
+
+    test03_btn.click(
+        fn=generate_prompter,
+        inputs=input_test03,
+        outputs=Textbox_test03
+    )
+
+    test04_btn.click(
+        fn=text_generate,
+        inputs=input_test04,
+        outputs=Textbox_test04
+    )
+    test05_btn.click(
+        fn=test05,
+        inputs=input_test05,
+        outputs=Textbox_test05
+    )
+    test06_btn.click(
+        fn=test06,
+        inputs=input_test06,
+        outputs=Textbox_test06
     )
 
 block.queue(max_size=64).launch(show_api=False, enable_queue=True, debug=True, share=False, server_name='0.0.0.0')

requirements.txt

@@ -1,5 +1,6 @@
 transformers==4.27.4
 torch==2.0.0
-gradio==3.24.1
+pytorch_lightning==1.7.7
+gradio==3.30.0
 sentencepiece==0.1.97
 sacremoses==0.0.53