app.py

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import spaces
import os

# bitsandbytesを無効化
os.environ["BITSANDBYTES_NOWELCOME"] = "1"

model_name = "minoD/JURAN"

# モデルのロード（CPUで、bitsandbytesを使わない）
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    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):
    result = f"""### 指示:あなたは企業の面接官です．就活生のエントリーシートを元に質問を行ってください．### 質問:{F}### 回答:"""
    result = result.replace('\n', '<NL>')
    return result

# テキスト生成関数の定義
@spaces.GPU(duration=60)  # タイムアウトを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():  # 勾配計算を無効化してメモリ節約
            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")
        else:
            return 'Warning: Expected prompt template to be emitted. Ignoring output.'
    
    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="半導体の研究に打ち込んだ"),
    outputs=gr.Textbox(label="想定される質問"),
    title="JURAN🌺",
    description="面接官モデルが回答を生成します。",
    api_name="ask",
    flagging_mode="never"
)

iface.launch(
    server_name="0.0.0.0",
    server_port=7860
)