Upload with huggingface_hub

app.py

@@ -0,0 +1,58 @@
+import gradio as gr
+import os
+from predict import *
+
+from transformers import T5ForConditionalGeneration
+from transformers import T5TokenizerFast as T5Tokenizer
+import pandas as pd
+model = "svjack/comet-atomic-zh"
+device = "cpu"
+#device = "cuda:0"
+tokenizer = T5Tokenizer.from_pretrained(model)
+model = T5ForConditionalGeneration.from_pretrained(model).to(device).eval()
+
+NEED_PREFIX = '以下事件有哪些必要的先决条件：'
+EFFECT_PREFIX = '下面的事件发生后可能会发生什么：'
+INTENT_PREFIX = '以下事件的动机是什么：'
+REACT_PREFIX = '以下事件发生后，你有什么感觉：'
+
+obj = Obj(model, tokenizer, device)
+
+text0 = "X吃到了一顿大餐。"
+text1 = "X和Y一起搭了个积木。"
+
+example_sample = [
+    [text0, False],
+    [text1, False],
+]
+
+def demo_func(event, do_sample):
+    #event = "X吃到了一顿大餐。"
+    times = 1
+    df = pd.DataFrame(
+    pd.Series(
+        [NEED_PREFIX, EFFECT_PREFIX, INTENT_PREFIX, REACT_PREFIX]
+    ).map(
+        lambda x: (x, [obj.predict(
+            "{}{}".format(x, event), do_sample = do_sample
+        )[0] for _ in range(times)][0])
+    ).values.tolist()
+    )
+    df.columns = ["PREFIX", "PRED"]
+    l = df.apply(lambda x: x.to_dict(), axis = 1).values.tolist()
+    return {
+        "Output": l
+    }
+
+demo = gr.Interface(
+        fn=demo_func,
+        inputs=[gr.Text(label = "Event"),
+                gr.Checkbox(label="do sample"),
+        ],
+        outputs="json",
+        title=f"Chinese Comet Atomic 🐰 demonstration",
+        examples=example_sample if example_sample else None,
+        cache_examples = False
+    )
+
+demo.launch(server_name=None, server_port=None)

predict.py

@@ -0,0 +1,47 @@
+class Obj:
+    def __init__(self, model, tokenizer, device = "cpu"):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = device
+        self.model = self.model.to(self.device)
+
+    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,
+            do_sample = do_sample
+        )
+        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

requirements.txt

@@ -0,0 +1,2 @@
+torch
+transformers