File size: 1,766 Bytes
6f93dce |
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 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
from flask import Flask, request, jsonify
from flask_cors import CORS
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
app = Flask(__name__)
CORS(app)
# Model configuration
MODEL_NAME = "deepseek-ai/deepseek-r1-6b-chat"
MAX_NEW_TOKENS = 512
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# Initialize model and tokenizer
try:
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
device_map="auto",
torch_dtype=torch.bfloat16 if DEVICE == "cuda" else torch.float32
)
print("Model loaded successfully!")
except Exception as e:
print(f"Model loading failed: {str(e)}")
model = None
def generate_response(prompt):
inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
outputs = model.generate(
**inputs,
max_new_tokens=MAX_NEW_TOKENS,
do_sample=True,
temperature=0.7,
top_p=0.9,
pad_token_id=tokenizer.eos_token_id
)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
@app.route('/chat', methods=['POST'])
def chat():
if not model:
return jsonify({"error": "Model not loaded"}), 500
data = request.json
prompt = data.get("prompt", "")
if not prompt:
return jsonify({"error": "No prompt provided"}), 400
try:
response = generate_response(prompt)
return jsonify({"response": response})
except Exception as e:
return jsonify({"error": str(e)}), 500
@app.route('/health', methods=['GET'])
def health_check():
status = "ready" if model else "unavailable"
return jsonify({"status": status})
if __name__ == '__main__':
app.run(host='0.0.0.0', port=5000) |