Update app.py

app.py

@@ -1,112 +1,133 @@
-import spaces
 import gradio as gr
-from transformers import AutoTokenizer, pipeline
+from huggingface_hub import InferenceClient
 import torch
+from transformers import AutoTokenizer, pipeline
 import logging
-import asyncio
-from functools import partial
+import spaces
 
-# ロギング設定
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+# ロガーの設定
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
 logger = logging.getLogger(__name__)
 
-# モデル定義
-classification_model_name = "unitary/toxic-bert"
-generation_model_name = "distilgpt2"  # 軽量なテキスト生成モデル
+# モデル定義（シンプルに2つだけ）
+TEXT_GENERATION_MODEL = "mistralai/Mistral-7B-v0.1"
+CLASSIFICATION_MODEL = "unitary/toxic-bert"
 
-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=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を利用する同期処理関数（分類）
-@spaces.GPU(duration=60)
-def classify_text(prompt):
-    logger.info(f"Running classification for: {prompt[:50]}...")
-    classification_result = classification_pipeline(prompt)
-    logger.info(f"Classification complete: {classification_result}")
-    return classification_result
+# モデルを事前ロード
+def load_text_generation_model():
+    logger.info(f"Loading text generation model: {TEXT_GENERATION_MODEL}")
+    tokenizer = AutoTokenizer.from_pretrained(TEXT_GENERATION_MODEL)
+    text_pipeline = pipeline(
+        "text-generation",
+        model=TEXT_GENERATION_MODEL,
+        tokenizer=tokenizer,
+        torch_dtype=torch.bfloat16,
+        trust_remote_code=True,
+        device_map="auto"
+    )
+    return text_pipeline
 
-# GPUを利用する同期処理関数（生成）
-@spaces.GPU(duration=60)
-def generate_text(prompt):
-    logger.info(f"Running text generation for: {prompt[:50]}...")
-    generation_result = generation_pipeline(
-        prompt,
-        max_new_tokens=50,
-        do_sample=True,
-        temperature=0.7,
-        num_return_sequences=1
+def load_classification_model():
+    logger.info(f"Loading classification model: {CLASSIFICATION_MODEL}")
+    tokenizer = AutoTokenizer.from_pretrained(CLASSIFICATION_MODEL)
+    class_pipeline = pipeline(
+        "text-classification",
+        model=CLASSIFICATION_MODEL,
+        tokenizer=tokenizer,
+        torch_dtype=torch.bfloat16,
+        trust_remote_code=True,
+        device_map="auto"
     )
-    generated_text = generation_result[0]["generated_text"]
-    logger.info(f"Text generation complete, generated: {len(generated_text)} chars")
-    return generated_text
+    return class_pipeline
 
-# 非同期ラッパー関数
-async def classify_text_async(prompt):
-    loop = asyncio.get_event_loop()
-    return await loop.run_in_executor(None, lambda: classify_text(prompt))
+# モデルのロード
+text_gen_pipeline = load_text_generation_model()
+classification_pipeline = load_classification_model()
 
-async def generate_text_async(prompt):
-    loop = asyncio.get_event_loop()
-    return await loop.run_in_executor(None, lambda: generate_text(prompt))
+# テキスト生成関数
+@spaces.GPU
+def generate_text(text):
+    try:
+        logger.info("Running text generation")
+        outputs = text_gen_pipeline(
+            text,
+            max_new_tokens=100,
+            do_sample=False,
+            num_return_sequences=1
+        )
+        return outputs[0]["generated_text"]
+    except Exception as e:
+        logger.error(f"Error in text generation: {str(e)}")
+        return f"Error: {str(e)}"
 
-# メイン処理を行う非同期関数
-async def process_text_async(prompt):
-    logger.info(f"Processing input asynchronously: {prompt[:50]}...")
-    
-    # 両方のタスクを並行して実行
-    classification_task = classify_text_async(prompt)
-    generation_task = generate_text_async(prompt)
-    
-    # 両方のタスクが完了するのを待つ
-    classification_result, generated_text = await asyncio.gather(
-        classification_task,
-        generation_task
-    )
-    
-    # 結果を組み合わせて返す
-    combined_result = f"分類結果: {classification_result}\n\n生成されたテキスト: {generated_text}"
-    return combined_result
+# テキスト分類関数
+@spaces.GPU
+def classify_text(text):
+    try:
+        logger.info("Running classification")
+        result = classification_pipeline(text)
+        return str(result)
+    except Exception as e:
+        logger.error(f"Error in classification: {str(e)}")
+        return f"Error: {str(e)}"
 
-# Gradio用の同期ラッパー関数
-def process_text(prompt):
-    # 非同期関数を同期的に実行
-    loop = asyncio.get_event_loop()
-    return loop.run_until_complete(process_text_async(prompt))
+# UIとロジックをシンプルに統合
+def create_ui():
+    with gr.Blocks() as demo:
+        gr.Markdown("""
+        # Toxic Eye (Simple Version)
+        This system evaluates the toxicity level of input text using two models.
+        """)
+        
+        with gr.Row():
+            input_text = gr.Textbox(
+                label="Input Text",
+                placeholder="Enter text to analyze...",
+                lines=3
+            )
+        
+        with gr.Row():
+            invoke_button = gr.Button(
+                "Analyze Text",
+                variant="primary",
+                size="lg"
+            )
+        
+        with gr.Tabs():
+            with gr.Tab("Text Generation"):
+                gen_output = gr.Textbox(
+                    label="Mistral-7B Output",
+                    lines=5,
+                    interactive=False
+                )
+            
+            with gr.Tab("Classification"):
+                class_output = gr.Textbox(
+                    label="Toxic-BERT Output",
+                    lines=5,
+                    interactive=False
+                )
+        
+        # イベントハンドラ
+        def handle_invoke(text):
+            gen_result = generate_text(text)
+            class_result = classify_text(text)
+            return gen_result, class_result
+        
+        invoke_button.click(
+            fn=handle_invoke,
+            inputs=[input_text],
+            outputs=[gen_output, class_output]
+        )
+    
+    return demo
 
-# Gradioインターフェース
-demo = gr.Interface(
-    fn=process_text,  # 同期ラッパー関数を使用
-    inputs=gr.Textbox(lines=3, label="入力テキスト"),
-    outputs=gr.Textbox(label="処理結果", lines=8),
-    title="テキスト分類 & 生成デモ (非同期処理版)",
-    description="入力テキストに対して分類と生成を非同期で並行実行します。"
-)
+def main():
+    demo = create_ui()
+    demo.launch()
 
-# アプリの起動
-logger.info("Starting application...")
-demo.launch()
+if __name__ == "__main__":
+    main()