Update app.py

app.py

@@ -1,31 +1,38 @@
+import spaces
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoTokenizer, pipeline
+import torch
 
-# 1. テキスト生成モデルの準備（DistilGPT2を使用）
-#    ZeroGPU環境ではGPUが無いのでCPUで動作。DistilGPT2は小型でCPUでも動作可能​:contentReference[oaicite:2]{index=2}。
+# 使用するモデル名（Falcon-7B-Instructを例に使用）
+model_name = "tiiuae/falcon-7b-instruct"
+
+# トークナイザとテキスト生成パイプラインの準備
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 generator = pipeline(
-    "text-generation", 
-    model="distilgpt2",           # 軽量な事前学習済みモデル
-    return_full_text=False        # プロンプト部分を含めず新たに生成したテキストのみ返す
+    "text-generation",
+    model=model_name,
+    tokenizer=tokenizer,
+    torch_dtype=torch.bfloat16,      # モデルをbfloat16精度でロード（メモリ節約）
+    trust_remote_code=True,          # モデルのリポジトリ内の追加コードを信頼して読み込む
+    device_map="auto"                # 利用可能なGPUに自動割り当て（ZeroGPU環境ではA100を使用）
 )
 
-# 2. ユーザー入力に対してテキスト生成を行う関数の定義
+# GPUを利用する推論関数を定義（ZeroGPUのためデコレータを使用）
+@spaces.GPU
 def generate_text(prompt):
-    # max_new_tokensで新規生成する単語数を制限（例: 最大50トークン）
-    result = generator(prompt, max_new_tokens=50)[0]["generated_text"]
-    return result
+    # プロンプトからテキストを生成し、結果文字列を返す
+    result = generator(prompt, max_new_tokens=100, do_sample=True)
+    generated = result[0]["generated_text"]
+    return generated
 
-# 3. Gradioインターフェイスの構築
-#    テキストボックスを入力と出力に設定し、generate_text関数を呼び出すシンプルなUIを構築
+# Gradioインタフェースの構築（テキスト入力→テキスト出力）
 demo = gr.Interface(
-    fn=generate_text, 
-    inputs=gr.components.Textbox(lines=4, label="入力プロンプト"), 
-    outputs=gr.components.Textbox(lines=5, label="生成テキスト"),
-    title="ZeroGPU対応: テキスト生成デモ",
-    description="プロンプトを入力すると、DistilGPT2モデルが続きを生成します。"
+    fn=generate_text,
+    inputs=gr.Textbox(lines=3, label="入力プロンプト"),
+    outputs=gr.Textbox(label="生成されたテキスト"),
+    title="Falcon-7B-Instruct テキスト生成デモ",
+    description="プロンプトを入力すると、大規模言語モデルが続きのテキストを生成します。"
 )
 
-# 4. アプリの起動（Spaces環境ではlaunch()不要だが、ローカル実行時には必要）
-if __name__ == "__main__":
-    demo.launch()
-    
+# アプリの起動（Spaces上ではこれによりサービスが公開される）
+demo.launch()