修改字詞的過濾

app.py

@@ -21,9 +21,9 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 big_processor = AutoProcessor.from_pretrained("microsoft/git-base-coco")
 big_model = AutoModelForCausalLM.from_pretrained("microsoft/git-base-coco")
 
-pipeline_01 = pipeline('text-generation', model='succinctly/text2image-prompt-generator')
-pipeline_02 = pipeline('text-generation', model='Gustavosta/MagicPrompt-Stable-Diffusion', tokenizer='gpt2')
-pipeline_03 = pipeline('text-generation', model='johnsu6616/ModelExport')
+pipeline_01 = pipeline('text-generation', model='succinctly/text2image-prompt-generator', max_new_tokens=256)
+pipeline_02 = pipeline('text-generation', model='Gustavosta/MagicPrompt-Stable-Diffusion', max_new_tokens=256)
+pipeline_03 = pipeline('text-generation', model='johnsu6616/ModelExport', max_new_tokens=256)
 
 zh2en_model = AutoModelForSeq2SeqLM.from_pretrained('Helsinki-NLP/opus-mt-zh-en').eval()
 zh2en_tokenizer = AutoTokenizer.from_pretrained('Helsinki-NLP/opus-mt-zh-en')
@@ -33,12 +33,15 @@ en2zh_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-zh")
 
 def translate_zh2en(text):
     with torch.no_grad():
-
         text = re.sub(r"[:\-–.!;?_#]", '', text)
+
         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', ',')
+
         text =re.sub(r'(?<![a-zA-Z])\s+|\s+(?![a-zA-Z])', '', text)
+
         text = re.sub(r',+', ',', text)
 
         encoded = zh2en_tokenizer([text], return_tensors='pt')
@@ -50,8 +53,12 @@ def translate_zh2en(text):
         if result == "No,no," :
             result = text
 
+        result = re.sub(r'<.*?>', '', result)
+
+        result = re.sub(r'\b(\w+)\b(?:\W+\1\b)+', r'\1', result, flags=re.IGNORECASE)
         return result
 
+
 def translate_en2zh(text):
     with torch.no_grad():
 
@@ -59,7 +66,7 @@ def translate_en2zh(text):
         sequences = en2zh_model.generate(**encoded)
         result = en2zh_tokenizer.batch_decode(sequences, skip_special_tokens=True)[0]
 
-        result = re.sub(r'(\b\w+\b)(?:\W+\1\b)+', r'\1', result)
+        result = re.sub(r'\b(\w+)\b(?:\W+\1\b)+', r'\1', result, flags=re.IGNORECASE)
         return result
 
 def load_prompter():
@@ -71,11 +78,12 @@ def load_prompter():
 
 prompter_model, prompter_tokenizer = load_prompter()
 
+
 def generate_prompter_pipeline_01(text):
     seed = random.randint(100, 1000000)
     set_seed(seed)
     text_in_english = translate_zh2en(text)
-    response = pipeline_01(text_in_english, max_new_tokens=80, num_return_sequences=3)
+    response = pipeline_01(text_in_english, num_return_sequences=3)
     response_list = []
     for x in response:
         resp = x['generated_text'].strip()
@@ -87,27 +95,27 @@ def generate_prompter_pipeline_01(text):
             response_list.append("\n")
 
     result = "".join(response_list)
-    result = re.sub('[^ ]+\.[^ ]+', '', result)
-    result = result.replace('<', '').replace('>', '')
+    result = re.sub('[^ ]+\.[^ ]+','', result)
+    result = result.replace("<", "").replace(">", "")
 
-    if result != '':
+    if result != "":
         return result
 
+
 def generate_prompter_tokenizer_01(text):
 
     text_in_english = translate_zh2en(text)
 
     input_ids = prompter_tokenizer(text_in_english.strip()+" Rephrase:", return_tensors="pt").input_ids
-
-    eos_id = 50256
+    
     outputs = prompter_model.generate(
         input_ids,
         do_sample=False,
-        max_new_tokens=80,
+
         num_beams=3,
         num_return_sequences=3,
-        pad_token_id=eos_id,
-        eos_token_id=eos_id,
+        pad_token_id= 50256,
+        eos_token_id = 50256,
         length_penalty=-1.0
     )
     output_texts = prompter_tokenizer.batch_decode(outputs, skip_special_tokens=True)
@@ -123,12 +131,11 @@ def generate_prompter_tokenizer_01(text):
         result.append("\n")
     return "".join(result)
 
-
 def generate_prompter_pipeline_02(text):
     seed = random.randint(100, 1000000)
     set_seed(seed)
     text_in_english = translate_zh2en(text)
-    response = pipeline_02(text_in_english, max_new_tokens=80, num_return_sequences=3)
+    response = pipeline_02(text_in_english, num_return_sequences=3)
     response_list = []
     for x in response:
         resp = x['generated_text'].strip()
@@ -149,13 +156,12 @@ def generate_prompter_pipeline_03(text):
     seed = random.randint(100, 1000000)
     set_seed(seed)
     text_in_english = translate_zh2en(text)
-    response = pipeline_03(text_in_english, max_new_tokens=80, num_return_sequences=3)
+    response = pipeline_03(text_in_english, num_return_sequences=3)
     response_list = []
     for x in response:
         resp = x['generated_text'].strip()
         if resp != text_in_english and len(resp) > (len(text_in_english) + 4):
 
-
             response_list.append(translate_en2zh(resp)+"\n")
             response_list.append(resp+"\n")
             response_list.append("\n")
@@ -184,7 +190,7 @@ def generate_render(text,choice):
 def get_prompt_from_image(input_image,choice):
     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_new_tokens=80)
+    generated_ids = big_model.to(device).generate(pixel_values=pixel_values)
     generated_caption = big_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
     text = re.sub(r"[:\-–.!;?_#]", '', generated_caption)
 
@@ -255,7 +261,7 @@ with gr.Blocks() as block:
         inputs=[input_text,radio_btn],
         outputs=[Textbox_1,Textbox_2]
       )
-    
+
     pic_prompter_btn.click(
         fn=get_prompt_from_image,
         inputs=[input_image,radio_btn],
@@ -292,6 +298,7 @@ with gr.Blocks() as block:
         outputs=Textbox_test05
     )
 
+
     test06_btn.click(
         fn= generate_prompter_pipeline_03,
         inputs= input_test06,