# Hugging FaceのSpaces側のコード
from transformers import AutoModelForCausalLM, AutoTokenizer
from flask import Flask, request, jsonify
import torch

app = Flask(__name__)

model_name = "rinna/youri-7b-chat"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

@app.route("/", methods=['POST'])
def endpoint():
    input_data = request.json
    input_text = input_data['input']

    # 推論処理
    inputs = tokenizer.encode(input_text, return_tensors='pt')
    outputs = model.generate(inputs, max_length=50, num_return_sequences=1)
    response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return jsonify({"output": response_text})

if __name__ == "__main__":
    app.run()