Update app.py

app.py

@@ -1,45 +1,33 @@
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import os
-import shutil
 import spaces
 
 model_name = "minoD/JURAN"
 
-# モデルのロード
+# モデルのロード（CPUで）
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     device_map="cpu",
     torch_dtype=torch.float16
-    
 )
-
 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
 def generate2(F=None, maxTokens=256):
+    # モデルをGPUに転送
     model.to("cuda")
+    
     # 推論
     prompt = generate_prompt(F)
-    input_ids = tokenizer(prompt,
-                          return_tensors="pt",
-                          truncation=True,
-                          add_special_tokens=False).input_ids.to(model.device)
+    input_ids = tokenizer(prompt, return_tensors="pt", truncation=True, add_special_tokens=False).input_ids.to("cuda")
     outputs = model.generate(
         input_ids=input_ids,
         max_new_tokens=maxTokens,
@@ -49,29 +37,31 @@ def generate2(F=None, maxTokens=256):
         top_k=40,
         no_repeat_ngram_size=2,
     )
+    
+    # CPUに戻す
+    model.to("cpu")
+    
     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>→改行
-        else:
-            return 'Warning: Expected prompt template to be emitted.  Ignoring output.'
-    else:
-       return 'Warning: no <eos> detected ignoring output'
     
-def inference(input_text):
-  return generate2(input_text)
+    # レスポンス内容のみ抽出
+    sentinel = "### 回答:"
+    sentinelLoc = decoded.find(sentinel)
+    if sentinelLoc >= 0:
+        result = decoded[sentinelLoc + len(sentinel):]
+        return result.replace("<NL>", "\n")
+    else:
+        return 'Warning: Expected prompt template to be emitted. Ignoring output.'
 
+def inference(input_text):
+    return generate2(input_text)
 
+# Gradioインターフェース
 iface = gr.Interface(
     fn=inference,
     inputs=gr.Textbox(lines=5, label="学生時代に打ち込んだこと、研究、ESを入力", placeholder="半導体の研究に打ち込んだ"),
@@ -79,8 +69,11 @@ iface = gr.Interface(
     title="JURAN🌺",
     description="面接官モデルが回答を生成します。",
     api_name="ask",
-    allow_flagging="never"
+    flagging_mode="never"
 )
 
-if __name__ == "__main__":
-    iface.launch()
+# if __name__ を削除して直接launch
+iface.launch(
+    server_name="0.0.0.0",
+    server_port=7860
+)