|
|
|
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("/app.py", methods=['POST']) |
|
def chat(): |
|
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() |
|
|
