Zmaker.pyをアップロード

前回のcommitでZmaker.pyを入れ忘れたので追加

Zmaker.py

@@ -0,0 +1,140 @@
+import torch
+from transformers import AutoModelForCausalLM, T5Tokenizer
+import csv, re, mojimoji
+
+class Zmaker:
+
+    #GPT2のモデル名
+    gpt_model_name = "rinna/japanese-gpt2-medium"
+
+    #文章の最大長
+    min_len, max_len = 1, 128
+
+    #予測時のパラメータ
+    top_k, top_p = 50, 0.95 #top-k検索の閾値
+    num_text = 1 #出力する文の数
+    temp = 1.0
+    repeat_ngram_size = 1
+
+    #推論にCPU利用を強制するか
+    use_cpu = True
+
+    def __init__(self, ft_path = None):
+        """コンストラクタ
+
+          コンストラクタ。モデルをファイルから読み込む場合と，
+          新規作成する場合で動作を分ける．
+
+          Args:
+              ft_path : ファインチューニングされたモデルのパス．Noneを指定すると
+              train : 学習を行うか
+          Returns:
+              なし
+        """
+
+        #モデルの設定
+        self.__SetModel(ft_path)
+
+        #モデルの状態をCPUかGPUかで切り替える
+        if self.use_cpu: #CPUの利用を強制する場合の処理
+            device = torch.device('cpu')
+        else: #特に指定が無いなら，GPUがあるときはGPUを使い，CPUのみの場合はCPUを使う
+            device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+        self.model.to(device)
+
+
+
+    def __SetModel(self, ft_path = None):
+      """GPT2の設定
+
+        GPT2のTokenizerおよびモデルを設定する.
+        ユーザー定義後と顔文字も語彙として認識されるように設定する．
+        
+        Args:
+            ft_path : ファインチューニング済みのモデルを読み込む
+                      何も指定しないとself.gpt_model_nameの事前学習モデルを
+                      ネットからダウンロードする．
+        Returns:
+            なし
+      """
+      #GPT2のTokenizerのインスタンスを生成
+      self.tokenizer = T5Tokenizer.from_pretrained(self.gpt_model_name)
+      self.tokenizer.do_lower_case = True # due to some bug of tokenizer config loading
+
+      #モデルの読み込み
+      if ft_path is not None:
+          self.model = AutoModelForCausalLM.from_pretrained(
+              ft_path, 
+              #torch_dtype = torch.bfloat16,
+              low_cpu_mem_usage = True
+          )
+      else:
+          print("fine-tuned model was not found")
+    
+      #モデルをevalモードに
+      self.model.eval()
+
+    def __TextCleaning(self, texts):
+        """テキストの前処理をする
+
+          テキストの前処理を行う．具体的に行うこととしては．．．
+          ・全角/半角スペースの除去
+          ・半角数字/アルファベットの全角化
+        """
+        #半角スペース，タブ，改行改ページを削除
+        texts = [re.sub("[\u3000 \t \s \n]", "", t) for t in texts]
+
+        #半角/全角を変換
+        texts = [mojimoji.han_to_zen(t) for t in texts]
+        return texts
+    
+
+    def GenLetter(self, prompt):
+      """怪文書の生成
+
+        GPT2で怪文書を生成する．
+        promptに続く文章を生成して出力する
+
+        Args:
+            prompt : 文章の先頭
+        Retunrs:
+            生成された文章のリスト
+      """
+
+      #テキストをクリーニング
+      prompt_clean = [prompt]
+      
+      #文章をtokenizerでエンコード
+      x = self.tokenizer.encode(
+          prompt_clean[0], return_tensors="pt", 
+          add_special_tokens=False
+      )
+      
+      #デバイスの選択
+      if self.use_cpu: #CPUの利用を強制する場合の処理
+          device = torch.device('cpu')
+      else: #特に指定が無いなら，GPUがあるときはGPUを使い，CPUのみの場合はCPUを使う
+          device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+      x = x.to(device)
+
+      #gpt2による推論
+      with torch.no_grad():
+          y = self.model.generate(
+              x, #入力
+              min_length=self.min_len,  # 文章の最小長
+              max_length=self.max_len,  # 文章の最大長
+              do_sample=True,   # 次の単語を確率で選ぶ
+              top_k=self.top_k, # Top-Kサンプリング
+              top_p=self.top_p,  # Top-pサンプリング
+              temperature=self.temp,  # 確率分布の調整
+              no_repeat_ngram_size = self.repeat_ngram_size, #同じ単語を何回繰り返していいか
+              num_return_sequences=self.num_text,  # 生成する文章の数
+              pad_token_id=self.tokenizer.pad_token_id,  # パディングのトークンID
+              bos_token_id=self.tokenizer.bos_token_id,  # テキスト先頭のトークンID
+              eos_token_id=self.tokenizer.eos_token_id,  # テキスト終端のトークンID
+              early_stopping=True
+          )
+      
+      # 特殊トークンをスキップして推論結果を文章にデコード
+      res = self.tokenizer.batch_decode(y, skip_special_tokens=True)
+      return res