Spaces:

ReikerXx
/

magda

Sleeping

magda / app.py

Update app.py

2831f7f verified 9 months ago

2.34 kB

from flask import Flask, request, jsonify
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# =========================
# 1️⃣ Load model & tokenizer
# =========================
model_name = "Qwen/Qwen2.5-0.5B-Instruct"

# Fast tokenizer for speed
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)

# Load model with correct dtype
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",  # Uses GPU if available, else CPU
    dtype=torch.float32  # CPU inference works better with float32
)

# Optional PyTorch 2.x compile (speeds up CPU inference)
if torch.__version__.startswith("2"):
    model = torch.compile(model)

# =========================
# 2️⃣ Hardcoded system prompt
# =========================
SYSTEM_PROMPT = "You are a friendly AI assistant that gives helpful and polite answers."

# =========================
# 3️⃣ Optimized chat function
# =========================
def chat(user_prompt: str):
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_prompt}
    ]

    # Apply Qwen chat template
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    # Encode input once
    inputs = tokenizer([text], return_tensors="pt").to(model.device)

    # Faster generation settings
    outputs = model.generate(
        **inputs,
        max_new_tokens=128,   # smaller = faster
        do_sample=False,      # deterministic = faster
        num_beams=1           # no beam search
    )

    # Decode only the first sequence
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# =========================
# 4️⃣ Flask app
# =========================
app = Flask(__name__)

@app.route("/chat", methods=["GET"])
def chat_route():
    user_message = request.args.get("message")

    if not user_message:
        return jsonify({"error": "No message provided"}), 400

    try:
        response = chat(user_message)
        return jsonify({"response": response})
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# =========================
# 5️⃣ Run Flask
# =========================
if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)

	from flask import Flask, request, jsonify
	from transformers import AutoModelForCausalLM, AutoTokenizer
	import torch

	# =========================
	# 1️⃣ Load model & tokenizer
	# =========================
	model_name = "Qwen/Qwen2.5-0.5B-Instruct"

	# Fast tokenizer for speed
	tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)

	# Load model with correct dtype
	model = AutoModelForCausalLM.from_pretrained(
	model_name,
	device_map="auto", # Uses GPU if available, else CPU
	dtype=torch.float32 # CPU inference works better with float32
	)

	# Optional PyTorch 2.x compile (speeds up CPU inference)
	if torch.__version__.startswith("2"):
	model = torch.compile(model)

	# =========================
	# 2️⃣ Hardcoded system prompt
	# =========================
	SYSTEM_PROMPT = "You are a friendly AI assistant that gives helpful and polite answers."

	# =========================
	# 3️⃣ Optimized chat function
	# =========================
	def chat(user_prompt: str):
	messages = [
	{"role": "system", "content": SYSTEM_PROMPT},
	{"role": "user", "content": user_prompt}
	]

	# Apply Qwen chat template
	text = tokenizer.apply_chat_template(
	messages,
	tokenize=False,
	add_generation_prompt=True
	)

	# Encode input once
	inputs = tokenizer([text], return_tensors="pt").to(model.device)

	# Faster generation settings
	outputs = model.generate(
	**inputs,
	max_new_tokens=128, # smaller = faster
	do_sample=False, # deterministic = faster
	num_beams=1 # no beam search
	)

	# Decode only the first sequence
	response = tokenizer.decode(outputs[0], skip_special_tokens=True)
	return response

	# =========================
	# 4️⃣ Flask app
	# =========================
	app = Flask(__name__)

	@app.route("/chat", methods=["GET"])
	def chat_route():
	user_message = request.args.get("message")

	if not user_message:
	return jsonify({"error": "No message provided"}), 400

	try:
	response = chat(user_message)
	return jsonify({"response": response})
	except Exception as e:
	return jsonify({"error": str(e)}), 500

	# =========================
	# 5️⃣ Run Flask
	# =========================
	if __name__ == "__main__":
	app.run(host="0.0.0.0", port=7860)