| |
| 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() |
|
|
