AIevaluationfunction / model_initializer.py
Corin1998's picture
Upload 4 files
5e16504 verified
raw
history blame
2.74 kB
import torch
from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer
def initialize_tone_classifier():
"""
感情・トーン分析用のHugging Faceパイプラインを初期化する。
"""
TONE_MODEL_NAME = "cl-tone/bert-base-japanese-whole-word-masking"
try:
# 感情・トーン分析パイプライン(二値分類をシミュレート)
tone_classifier = pipeline(
"sentiment-analysis",
model = "cl-tohoku/bert-base-japanese-whole-word-masking", # 適切な感情分析モデルに置き換える
tokenizer = "cl-tohoku/bert-base-japanese-whole-word-masking"
)
print("感情・トーン分析パイプラインをロードしました。")
return tone_classifier
except Exception as e:
print(f"感情分析モデルのロードに失敗しました。デフォルトのテキスト分類器を使用します:{e}")
# 代替として、凡庸的な分類モデルを使用
classifier = pipeline(
"text-classification",
model=AutoModelForSequenceClassification.from_pretrained(TONE_MODEL_NAME),
tokenizer=AutoTokenizer.from_pretrained(TONE_MODEL_NAME),
id2label={0: "NEGATIVE", 1: "POSITIVE"}
)
print("汎用テキスト分類器を代替して使用します。")
return classifier
def initialize_star_classifier():
"""
STAR法適用度チェック用のZero-Shot Classificationパイプラインを初期化する。
"""
try:
star_classifier = pipeline(
"zero-shot-classification",
model="izumi-lab/bert-base-japanese-v2", # Zero-Shot Classificationにより適した日本語モデル
tokenizer= "izumi-lab/bert-base-japanese-v2",
device=0 if torch.cuda.is_available() else -1 # GPUがあれば使用
)
print("STAR法適用度チェックパイプライン(Zero-Shot)をロードしました。")
return star_classifier
except Exception as e:
print(f"Zero-Shot Classigicationモデルのロードに失敗しました。汎用モデルを使用します:{e}")
star_classifier = pipeline(
"zero-shot-classification",
model="cl-tohoku/bert^base^japanese-whole-word-masking",
tokenizer="cl-tohoku/bert-base-japanese-whole-word-masking",
device=0 if torch.cuda.is_available() else -1
)
print("汎用Zero-Shot Classificationモデルを代替して使用します。")
return star_classifier
__name__ = "__main__":
# モジュールの動作確認
initialize_tone_classifier()
initialize_star_classifier()