Update app.py

app.py

@@ -20,22 +20,26 @@ tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
 
 # プロンプトテンプレートの準備
 def generate_prompt(F):
-    result = f"""### 指示:あなたは企業の面接官です．就活生のエントリーシートを元に質問を行ってください．### 質問:{F}### 回答:"""
+    result = f"""### 指示:
+あなたは企業の面接官です。以下の就活生のエントリーシート内容を読んで、深掘りする質問を1つ考えてください。
+
+### エントリーシート:
+{F}
+
+### 面接官の質問:"""
     result = result.replace('\n', '<NL>')
     return result
 
 # テキスト生成関数の定義
-@spaces.GPU(duration=60)  # タイムアウトを60秒に設定
+@spaces.GPU(duration=60)
 def generate2(F=None, maxTokens=256):
     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():  # 勾配計算を無効化してメモリ節約
+        with torch.no_grad():
             outputs = model.generate(
                 input_ids=input_ids,
                 max_new_tokens=maxTokens,
@@ -46,9 +50,8 @@ def generate2(F=None, maxTokens=256):
                 no_repeat_ngram_size=2,
             )
         
-        # CPUに戻す
         model.to("cpu")
-        torch.cuda.empty_cache()  # GPUメモリをクリア
+        torch.cuda.empty_cache()
         
         outputs = outputs[0].tolist()
         decoded = tokenizer.decode(outputs)
@@ -58,12 +61,14 @@ def generate2(F=None, maxTokens=256):
             eos_index = outputs.index(tokenizer.eos_token_id)
             decoded = tokenizer.decode(outputs[:eos_index])
         
-        # レスポンス内容のみ抽出
-        sentinel = "### 回答:"
+        # レスポンス内容のみ抽出（修正）
+        sentinel = "### 面接官の質問:"
         sentinelLoc = decoded.find(sentinel)
         if sentinelLoc >= 0:
             result = decoded[sentinelLoc + len(sentinel):]
-            return result.replace("<NL>", "\n")
+            # 最初の改行までを取得（1つの質問だけ）
+            result = result.split('\n')[0] if '\n' in result else result
+            return result.replace("<NL>", "\n").strip()
         else:
             return 'Warning: Expected prompt template to be emitted. Ignoring output.'