Update app.py

app.py

@@ -1,20 +1,22 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from transformers import pipeline
 import os
 from typing import Optional
 
-# FastAPIアプリの初期化
-app = FastAPI(
-    title="Hugging Face API on Spaces",
-    description="Hugging Face Transformersを使ったAPI",
-    version="1.0.0"
-)
+# キャッシュディレクトリの設定
+os.environ["HF_HOME"] = "/app/.cache"
+os.environ["TRANSFORMERS_CACHE"] = "/app/.cache"
+os.environ["HUGGINGFACE_HUB_CACHE"] = "/app/.cache"
+
+# キャッシュディレクトリを作成
+cache_dir = "/app/.cache"
+os.makedirs(cache_dir, exist_ok=True)
+
+app = FastAPI(title="Lightweight Hugging Face API")
 
-# リクエスト用のPydanticモデル
 class TextRequest(BaseModel):
     text: str
-    max_length: Optional[int] = 100
+    max_length: Optional[int] = 50
 
 class SentimentResponse(BaseModel):
     text: str
@@ -27,62 +29,59 @@ class GenerateResponse(BaseModel):
     generated_text: str
     model_name: str
 
-# グローバル変数でモデルを保持
+# グローバル変数
 sentiment_classifier = None
 text_generator = None
 
 @app.on_event("startup")
 async def load_models():
-    """アプリ起動時にモデルをロード"""
+    """軽量モデルをロード"""
     global sentiment_classifier, text_generator
     
-    print("モデルをロード中...")
+    print("🚀 軽量モデルのロード開始...")
     
     try:
-        # 感情分析モデル（軽量版を使用）
-        print("Loading sentiment analysis model...")
+        from transformers import pipeline
+        
+        # 非常に軽量な感情分析モデル
+        print("📥 軽量感情分析モデルをロード中...")
         sentiment_classifier = pipeline(
             "sentiment-analysis",
-            model="cardiffnlp/twitter-roberta-base-sentiment-latest",
-            return_all_scores=False
+            model="prajjwal1/bert-tiny",  # 非常に軽量
+            cache_dir="/app/.cache"
         )
-        print("✅ Sentiment model loaded")
+        print("✅ 感情分析モデル読み込み完了")
         
-        # テキスト生成モデル（軽量版）
-        print("Loading text generation model...")
+        # 軽量テキスト生成モデル
+        print("📥 軽量テキスト生成モデルをロード中...")
         text_generator = pipeline(
             "text-generation",
-            model="distilgpt2",
-            pad_token_id=50256  # GPT-2のEOSトークンID
+            model="sshleifer/tiny-gpt2",  # 非常に軽量
+            cache_dir="/app/.cache"
         )
-        print("✅ Text generation model loaded")
+        print("✅ テキスト生成モデル読み込み完了")
         
-        print("✅ 全てのモデルのロードが完了しました")
+        print("✅ 全てのモデル読み込み完了")
         
     except Exception as e:
-        print(f"❌ モデルロードエラー: {e}")
-        print(f"エラーの詳細: {type(e).__name__}")
+        print(f"❌ モデル読み込みエラー: {e}")
         import traceback
         traceback.print_exc()
 
 @app.get("/")
 async def root():
-    """ヘルスチェック用エンドポイント"""
     return {
-        "message": "🤗 Hugging Face API is running on Spaces!",
+        "message": "🤗 Lightweight Hugging Face API is running!",
         "status": "healthy",
-        "endpoints": ["/sentiment", "/generate", "/models"]
+        "models": "lightweight versions"
     }
 
 @app.post("/sentiment", response_model=SentimentResponse)
 async def analyze_sentiment(request: TextRequest):
-    """感情分析エンドポイント"""
+    """軽量感情分析"""
     try:
         if sentiment_classifier is None:
-            raise HTTPException(
-                status_code=503, 
-                detail="Sentiment model not loaded. Please try again later."
-            )
+            raise HTTPException(status_code=503, detail="Sentiment model not loaded")
         
         result = sentiment_classifier(request.text)
         
@@ -90,7 +89,7 @@ async def analyze_sentiment(request: TextRequest):
             text=request.text,
             sentiment=result[0]["label"],
             confidence=round(result[0]["score"], 4),
-            model_name="cardiffnlp/twitter-roberta-base-sentiment-latest"
+            model_name="prajjwal1/bert-tiny"
         )
     
     except Exception as e:
@@ -98,98 +97,46 @@ async def analyze_sentiment(request: TextRequest):
 
 @app.post("/generate", response_model=GenerateResponse)
 async def generate_text(request: TextRequest):
-    """テキスト生成エンドポイント"""
+    """軽量テキスト生成"""
     try:
         if text_generator is None:
-            raise HTTPException(
-                status_code=503, 
-                detail="Text generation model not loaded. Please try again later."
-            )
+            raise HTTPException(status_code=503, detail="Generation model not loaded")
         
-        # 入力テキストの検証
-        if not request.text or len(request.text.strip()) == 0:
-            raise HTTPException(status_code=400, detail="Text cannot be empty")
-        
-        # Spacesの制限を考慮して短めに設定
         max_length = min(request.max_length, 100)
-        input_length = len(request.text.split())
         
-        # 入力より長い出力を生成するように調整
-        if max_length <= input_length:
-            max_length = input_length + 20
+        result = text_generator(
+            request.text,
+            max_length=max_length,
+            num_return_sequences=1,
+            temperature=0.7,
+            do_sample=True,
+            pad_token_id=text_generator.tokenizer.eos_token_id
+        )
         
-        try:
-            result = text_generator(
-                request.text,
-                max_length=max_length,
-                num_return_sequences=1,
-                temperature=0.7,
-                do_sample=True,
-                truncation=True,
-                pad_token_id=text_generator.tokenizer.eos_token_id
-            )
-            
-            generated_text = result[0]["generated_text"]
-            
-            return GenerateResponse(
-                input_text=request.text,
-                generated_text=generated_text,
-                model_name="distilgpt2"
-            )
-            
-        except Exception as model_error:
-            # モデル固有のエラーをキャッチ
-            print(f"Model error: {model_error}")
-            raise HTTPException(
-                status_code=500, 
-                detail=f"Model processing failed: {str(model_error)}"
-            )
+        return GenerateResponse(
+            input_text=request.text,
+            generated_text=result[0]["generated_text"],
+            model_name="sshleifer/tiny-gpt2"
+        )
     
-    except HTTPException:
-        # HTTPExceptionは再発生
-        raise
     except Exception as e:
-        print(f"Unexpected error: {e}")
         raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")
 
 @app.get("/models")
 async def get_models():
-    """利用可能なモデル情報を取得"""
     return {
         "sentiment_analysis": {
-            "model": "cardiffnlp/twitter-roberta-base-sentiment-latest",
+            "model": "prajjwal1/bert-tiny",
             "status": "loaded" if sentiment_classifier else "not loaded"
         },
         "text_generation": {
-            "model": "distilgpt2",
+            "model": "sshleifer/tiny-gpt2", 
             "status": "loaded" if text_generator else "not loaded"
         },
-        "platform": "Hugging Face Spaces"
-    }
-
-@app.get("/debug")
-async def debug_info():
-    """デバッグ情報を取得"""
-    import sys
-    import torch
-    
-    return {
-        "python_version": sys.version,
-        "torch_version": torch.__version__ if 'torch' in sys.modules else "not installed",
-        "models_loaded": {
-            "sentiment": sentiment_classifier is not None,
-            "generator": text_generator is not None
-        },
-        "generator_tokenizer": {
-            "vocab_size": text_generator.tokenizer.vocab_size if text_generator else None,
-            "eos_token_id": text_generator.tokenizer.eos_token_id if text_generator else None,
-            "pad_token_id": text_generator.tokenizer.pad_token_id if text_generator else None
-        } if text_generator else None
+        "note": "Using lightweight models for Spaces compatibility"
     }
 
-# Spaces用の追加設定
 if __name__ == "__main__":
     import uvicorn
-    # Spacesでは通常ポート7860を使用
     port = int(os.environ.get("PORT", 7860))
     uvicorn.run(app, host="0.0.0.0", port=port)