File size: 1,183 Bytes
7339449 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
import flask
from flask import request, jsonify
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
app = flask.Flask(__name__)
model_id = "facebook/blenderbot-400M-distill"
print("🔄 Loading fast chat model...")
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
print("✅ Model loaded instantly!")
@app.route('/chat', methods=['POST'])
def chat():
try:
data = request.get_json()
msg = data.get("message", "")
if not msg:
return jsonify({"error": "No message sent"}), 400
inputs = tokenizer(msg, return_tensors="pt").to(device)
output = model.generate(
**inputs,
max_length=200,
do_sample=True,
top_p=0.92,
temperature=0.7
)
reply = tokenizer.decode(output[0], skip_special_tokens=True)
return jsonify({"reply": reply})
except Exception as e:
return jsonify({"error": str(e)}), 500
if __name__ == "__main__":
app.run(host='0.0.0.0', port=7860) |