Update app.py

app.py

@@ -4,12 +4,12 @@ from transformers import AutoTokenizer, pipeline
 import torch
 
 #model_name = "tiiuae/falcon-7b-instruct"
-model_name = "meta-llama/Llama-2-7b-hf"
+model_name = "unitary/toxic-bert"
 
 # トークナイザとテキスト生成パイプラインの準備
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 generator = pipeline(
-    "text-generation",
+    "text-classification",
     model=model_name,
     tokenizer=tokenizer,
     torch_dtype=torch.bfloat16,      # モデルをbfloat16精度でロード（メモリ節約）
@@ -22,7 +22,9 @@ generator = pipeline(
 def generate_text(prompt):
     # プロンプトからテキストを生成し、結果文字列を返す
     result = generator(prompt, max_new_tokens=100, do_sample=True)
-    generated = result[0]["generated_text"]
+    #generated = result[0]["generated_text"]
+    generated = result
+    
     return generated
 
 # Gradioインタフェースの構築（テキスト入力→テキスト出力）