main.py

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

local_path = "./models/roberta-large"

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)
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    os.makedirs(local_path, exist_ok=True)
    model.save_pretrained(local_path)
    tokenizer.save_pretrained(local_path)

app = Flask(__name__)


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


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