normalpurpose3/Ignite-Pro-Workspace

1import os2import torch3from flask import Flask, request, jsonify4from transformers import AutoModelForCausalLM, AutoTokenizer5 6app = Flask(__name__)7 8# --- CONFIGURATION ---9# Use CPU to ensure it works on the Free Tier without NVIDIA errors10DEVICE = "cpu"11MODEL_PATH = "./my_dream_model" 12BASE_MODEL = "google/gemma-3-4b-it"13TOKEN = os.environ.get("HF_TOKEN")14 15# Global variables to store the brain16tokenizer = None17model = None18 19print("System: Loading Version 1 Engine...")20 21try:22    # Load tokenizer and model specifically for CPU (float32)23    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, token=TOKEN)24    model = AutoModelForCausalLM.from_pretrained(25        MODEL_PATH, 26        torch_dtype=torch.float32, 27        device_map={"": DEVICE}28    )29    print("System: Version 1 is ONLINE.")30except Exception as e:31    print(f"Startup Error: {e}")32 33@app.route("/")34def index():35    # Serves your Red-to-Orange UI36    with open("index.html", "r") as f:37        return f.read()38 39@app.route("/chat", methods=["POST"])40def chat():41    global model, tokenizer42    43    try:44        # 1. Check if model is ready45        if model is None or tokenizer is None:46            return jsonify({"reply": "AI is still booting up... please wait a moment."})47 48        # 2. Get user input49        data = request.json50        user_input = data.get("message", "")51 52        # 3. Simple, effective prompt53        prompt = f"<bos><start_of_turn>user\n{user_input}<end_of_turn>\n<start_of_turn>model\n"54        55        # 4. Process56        inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)57        58        with torch.no_grad():59            outputs = model.generate(60                **inputs, 61                max_new_tokens=200, 62                do_sample=True,63                temperature=0.7,64                pad_token_id=tokenizer.eos_token_id65            )66        67        # 5. Clean and send reply68        full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)69        reply = full_text.split("model")[-1].strip()70 71        return jsonify({"reply": reply})72 73    except Exception as e:74        return jsonify({"reply": f"Error: {str(e)}"})75 76if __name__ == "__main__":77    # Hugging Face Spaces standard port78    app.run(host="0.0.0.0", port=7860)