main.py

from flask import Flask, request, jsonify
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
import os

local_path = "./models/roberta-large"

os.environ["HF_HOME"] = "/app/hf_home"
os.environ["TRANSFORMERS_CACHE"] = "/app/cache"  # 병행 사용 가능

model_id = "klue/roberta-large"

if os.path.exists(local_path):
    print("🔄 모델 로컬에서 로드 중...")
    model = AutoModelForSequenceClassification.from_pretrained(local_path)
    tokenizer = AutoTokenizer.from_pretrained(local_path)
else:
    print("⬇️ 모델 허깅페이스에서 다운로드 중...")
    model = AutoModelForSequenceClassification.from_pretrained(
        model_id, cache_dir=os.environ["HF_HOME"]
    )
    tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=os.environ["HF_HOME"])
    os.makedirs(local_path, exist_ok=True)
    model.save_pretrained(local_path)
    tokenizer.save_pretrained(local_path)

app = Flask(__name__)

print("🔄 모델 로드 완료")


@app.route("/generate", methods=["GET"])
def generate():
    return jsonify({"result": "generate/get"})


@app.route("/generate", methods=["POST"])
def generate_post():
    data = request.json
    print(data)
    return jsonify({"result": "generate/post"})


@app.route("/", methods=["GET"])
def index():
    return jsonify({"result": "success"})


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)