Update app.py

app.py

@@ -2,39 +2,77 @@ import spaces
 import gradio as gr
 from transformers import AutoTokenizer, pipeline
 import torch
+import logging
 
-#model_name = "tiiuae/falcon-7b-instruct"
-model_name = "unitary/toxic-bert"
+# ロギング設定
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
-# トークナイザとテキスト生成パイプラインの準備
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-generator = pipeline(
+# モデル定義
+classification_model_name = "unitary/toxic-bert"
+generation_model_name = "distilgpt2"  # 軽量なテキスト生成モデル
+
+logger.info("Starting model loading...")
+
+# 分類モデルのロード
+logger.info(f"Loading classification model: {classification_model_name}")
+classification_tokenizer = AutoTokenizer.from_pretrained(classification_model_name)
+classification_pipeline = pipeline(
     "text-classification",
-    model=model_name,
-    tokenizer=tokenizer,
-    torch_dtype=torch.bfloat16,      # モデルをbfloat16精度でロード（メモリ節約）
-    trust_remote_code=True,          # モデルのリポジトリ内の追加コードを信頼して読み込む
-    device_map="auto"                # 利用可能なGPUに自動割り当て（ZeroGPU環境ではA100を使用）
+    model=classification_model_name,
+    tokenizer=classification_tokenizer,
+    torch_dtype=torch.bfloat16,
+    trust_remote_code=True,
+    device_map="auto"
+)
+logger.info(f"Classification model loaded successfully: {classification_model_name}")
+
+# 生成モデルのロード
+logger.info(f"Loading generation model: {generation_model_name}")
+generation_tokenizer = AutoTokenizer.from_pretrained(generation_model_name)
+generation_pipeline = pipeline(
+    "text-generation",
+    model=generation_model_name,
+    tokenizer=generation_tokenizer,
+    torch_dtype=torch.bfloat16,
+    trust_remote_code=True,
+    device_map="auto"
 )
+logger.info(f"Generation model loaded successfully: {generation_model_name}")
 
-# GPUを利用する推論関数を定義（ZeroGPUのためデコレータを使用）
+# GPUを利用する推論関数（両方のモデルを使用）
 @spaces.GPU(duration=120)
-def generate_text(prompt):
-    # プロンプトからテキストを生成し、結果文字列を返す
-    result = generator(prompt)
-    #generated = result[0]["generated_text"]
-    generated = result
+def process_text(prompt):
+    logger.info(f"Processing input: {prompt[:50]}...")
+    
+    # 分類モデルで処理
+    classification_result = classification_pipeline(prompt)
+    logger.info(f"Classification complete: {classification_result}")
+    
+    # 生成モデルで処理
+    generation_result = generation_pipeline(
+        prompt,
+        max_new_tokens=50,
+        do_sample=True,
+        temperature=0.7,
+        num_return_sequences=1
+    )
+    generated_text = generation_result[0]["generated_text"]
+    logger.info(f"Text generation complete, generated: {len(generated_text)} chars")
     
-    return generated
+    # 結果を組み合わせて返す
+    combined_result = f"分類結果: {classification_result}\n\n生成されたテキスト: {generated_text}"
+    return combined_result
 
-# Gradioインタフェースの構築（テキスト入力→テキスト出力）
+# Gradioインタフェース
 demo = gr.Interface(
-    fn=generate_text,
-    inputs=gr.Textbox(lines=3, label="入力プロンプト"),
-    outputs=gr.Textbox(label="生成されたテキスト"),
-    title="Falcon-7B-Instruct テキスト生成デモ",
-    description="プロンプトを入力すると、大規模言語モデルが続きのテキストを生成します。"
+    fn=process_text,
+    inputs=gr.Textbox(lines=3, label="入力テキスト"),
+    outputs=gr.Textbox(label="処理結果", lines=8),
+    title="テキスト分類 & 生成デモ",
+    description="入力テキストに対して分類と生成の両方を実行します。"
 )
 
-# アプリの起動（Spaces上ではこれによりサービスが公開される）
+# アプリの起動
+logger.info("Starting application...")
 demo.launch()