Upload 3 files

predict.py

@@ -0,0 +1,59 @@
+import re
+
+def batch_as_list(a, batch_size = int(100000)):
+    req = []
+    for ele in a:
+        if not req:
+            req.append([])
+        if len(req[-1]) < batch_size:
+            req[-1].append(ele)
+        else:
+            req.append([])
+            req[-1].append(ele)
+    return req
+
+class Obj:
+    def __init__(self, model, tokenizer, device = "cpu"):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = "cpu"
+
+    def predict(
+        self,
+        source_text: str,
+        max_length: int = 512,
+        num_return_sequences: int = 1,
+        num_beams: int = 2,
+        top_k: int = 50,
+        top_p: float = 0.95,
+        do_sample: bool = True,
+        repetition_penalty: float = 2.5,
+        length_penalty: float = 1.0,
+        early_stopping: bool = True,
+        skip_special_tokens: bool = True,
+        clean_up_tokenization_spaces: bool = True,
+    ):
+        input_ids = self.tokenizer.encode(
+            source_text, return_tensors="pt", add_special_tokens=True
+        )
+        input_ids = input_ids.to(self.device)
+        generated_ids = self.model.generate(
+            input_ids=input_ids,
+            num_beams=num_beams,
+            max_length=max_length,
+            repetition_penalty=repetition_penalty,
+            length_penalty=length_penalty,
+            early_stopping=early_stopping,
+            top_p=top_p,
+            top_k=top_k,
+            num_return_sequences=num_return_sequences,
+        )
+        preds = [
+            self.tokenizer.decode(
+                g,
+                skip_special_tokens=skip_special_tokens,
+                clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            )
+            for g in generated_ids
+        ]
+        return preds

reconstructor.py

@@ -0,0 +1,39 @@
+from predict import *
+from transformers import (
+    T5ForConditionalGeneration,
+    T5TokenizerFast as T5Tokenizer,
+)
+import jieba.posseg as posseg
+
+model_path = "svjack/T5-dialogue-collect-v5"
+tokenizer = T5Tokenizer.from_pretrained(model_path)
+model = T5ForConditionalGeneration.from_pretrained(model_path)
+
+rec_obj = Obj(model, tokenizer)
+
+def process_one_sent(input_):
+    assert type(input_) == type("")
+    input_ = " ".join(map(lambda y: y.word.strip() ,filter(lambda x: x.flag != "x" ,
+    posseg.lcut(input_))))
+    return input_
+
+def predict_split(sp_list, cut_tokens = True):
+    assert type(sp_list) == type([])
+    if cut_tokens:
+        src_text = '''
+            根据下面的上下文进行分段：
+            上下文：{}
+            答案：
+            '''.format(" ".join(
+            map(process_one_sent ,sp_list)
+            ))
+    else:
+        src_text = '''
+            根据下面的上下文进行分段：
+            上下文：{}
+            答案：
+            '''.format("".join(sp_list))
+    print(src_text)
+    pred = rec_obj.predict(src_text)[0]
+    pred = list(filter(lambda y: y ,map(lambda x: x.strip() ,pred.split("分段:"))))
+    return pred

requirements.txt

@@ -0,0 +1,3 @@
+transformers==4.20.1
+jieba
+gradio