Update app.py

app.py

@@ -1,74 +1,79 @@
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import spaces
 import os
-import shutil
+
+# bitsandbytesを無効化
+os.environ["BITSANDBYTES_NOWELCOME"] = "1"
 
 model_name = "minoD/JURAN"
 
-# モデルのロード
+# モデルのロード（CPUで、bitsandbytesを使わない）
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
-    device_map="auto",
-    torch_dtype=torch.float16
+    device_map="cpu",
+    torch_dtype=torch.float16,
+    low_cpu_mem_usage=True,  # メモリ効率を改善
 )
-
 tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
 
 # プロンプトテンプレートの準備
 def generate_prompt(F):
-    # input キーの代わりに Q と F を使用
-    result = f"""### 指示:あなたは企業の面接官です．就活生のエントリーシートを元に質問を行ってください．
-
-### 質問:
-{F}
-
-### 回答:
-""" # 回答セクションを追加
-    # 改行→<NL>
+    result = f"""### 指示:あなたは企業の面接官です．就活生のエントリーシートを元に質問を行ってください．### 質問:{F}### 回答:"""
     result = result.replace('\n', '<NL>')
     return result
 
 # テキスト生成関数の定義
+@spaces.GPU(duration=60)  # タイムアウトを60秒に設定
 def generate2(F=None, maxTokens=256):
-    # 推論
-    prompt = generate_prompt(F)
-    input_ids = tokenizer(prompt,
-                          return_tensors="pt",
-                          truncation=True,
-                          add_special_tokens=False).input_ids.to(model.device)
-    outputs = model.generate(
-        input_ids=input_ids,
-        max_new_tokens=maxTokens,
-        do_sample=True,
-        temperature=0.7,
-        top_p=0.75,
-        top_k=40,
-        no_repeat_ngram_size=2,
-    )
-    outputs = outputs[0].tolist()
-    decoded = tokenizer.decode(outputs)
-
-    # EOSトークンにヒットしたらデコード完了
-    if tokenizer.eos_token_id in outputs:
-        eos_index = outputs.index(tokenizer.eos_token_id)
-        decoded = tokenizer.decode(outputs[:eos_index])
-
+    try:
+        # モデルをGPUに転送
+        model.to("cuda")
+        
+        # 推論
+        prompt = generate_prompt(F)
+        input_ids = tokenizer(prompt, return_tensors="pt", truncation=True, add_special_tokens=False).input_ids.to("cuda")
+        
+        with torch.no_grad():  # 勾配計算を無効化してメモリ節約
+            outputs = model.generate(
+                input_ids=input_ids,
+                max_new_tokens=maxTokens,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.75,
+                top_k=40,
+                no_repeat_ngram_size=2,
+            )
+        
+        # CPUに戻す
+        model.to("cpu")
+        torch.cuda.empty_cache()  # GPUメモリをクリア
+        
+        outputs = outputs[0].tolist()
+        decoded = tokenizer.decode(outputs)
+        
+        # EOSトークンにヒットしたらデコード完了
+        if tokenizer.eos_token_id in outputs:
+            eos_index = outputs.index(tokenizer.eos_token_id)
+            decoded = tokenizer.decode(outputs[:eos_index])
+        
         # レスポンス内容のみ抽出
         sentinel = "### 回答:"
         sentinelLoc = decoded.find(sentinel)
         if sentinelLoc >= 0:
             result = decoded[sentinelLoc + len(sentinel):]
-            return result.replace("<NL>", "\n")  # <NL>→改行
+            return result.replace("<NL>", "\n")
         else:
-            return 'Warning: Expected prompt template to be emitted.  Ignoring output.'
-    else:
-       return 'Warning: no <eos> detected ignoring output'
+            return 'Warning: Expected prompt template to be emitted. Ignoring output.'
     
-def inference(input_text):
-  return generate2(input_text)
+    except Exception as e:
+        return f"エラーが発生しました: {str(e)}"
 
+def inference(input_text):
+    return generate2(input_text)
 
+# Gradioインターフェース
 iface = gr.Interface(
     fn=inference,
     inputs=gr.Textbox(lines=5, label="学生時代に打ち込んだこと、研究、ESを入力", placeholder="半導体の研究に打ち込んだ"),
@@ -76,8 +81,10 @@ iface = gr.Interface(
     title="JURAN🌺",
     description="面接官モデルが回答を生成します。",
     api_name="ask",
-    allow_flagging="never"
+    flagging_mode="never"
 )
 
-if __name__ == "__main__":
-     iface.launch(share=True)
+iface.launch(
+    server_name="0.0.0.0",
+    server_port=7860
+)