Update app.py

app.py

@@ -1,70 +1,61 @@
 from flask import Flask, request, jsonify
 from flask_cors import CORS
 import os
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+import requests
 
 app = Flask(__name__)
-CORS(app)
+CORS(app)  # Keeps your Acode terminal connected securely!
 
-MODEL_NAME = "unsloth/mistral-7b-instruct-v0.3-bnb-4bit"
+# Pulls your hidden WRITE token from Space secrets
 HF_TOKEN = os.getenv("HF_TOKEN")
-
-print("Initializing local memory nodes... Downloading Unsloth bits...")
-
-# 1. Load the tokenizer matching your exact model structure
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_TOKEN)
-
-# 2. Configure 4-bit loading directly inside your Space container memory
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.float16
-)
-
-# 3. Pull weights directly into local cache space
-try:
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        quantization_config=bnb_config,
-        device_map="auto",
-        token=HF_TOKEN
-    )
-    print("Krypton-1 Core Engine Fully Loaded in Space Memory.")
-except Exception as e:
-    print(f"Boot Failure Error: {str(e)}")
-    model = None
+MODEL_REPO = "unsloth/mistral-7b-instruct-v0.3-bnb-4bit"
+API_URL = f"https://api-inference.huggingface.co/models/{MODEL_REPO}"
 
 @app.route("/", methods=["GET"])
 def home():
-    return "Krypton-1 Dedicated Unsloth Node is Online."
+    return "Krypton-1 Core Unsloth Node Active."
 
 @app.route("/api/chat", methods=["POST"])
 def chat():
-    if model is None:
-        return jsonify({"reply": "SYS_ERR: Engine failed to cache local weights. Check container storage logs."})
-        
     try:
         data = request.json
         user_prompt = data.get("prompt", "")
         
-        # Structure the prompt properly for Mistral Instruct layout
-        messages = [{"role": "user", "content": user_prompt}]
-        inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
-        
-        # Generate raw response tokens straight from your loaded weights
-        outputs = model.generate(inputs, max_new_tokens=250, temperature=0.7, do_sample=True)
+        headers = {
+            "Authorization": f"Bearer {HF_TOKEN}",
+            "Content-Type": "application/json"
+        }
         
-        # Decode the output text cleanly
-        decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # Structure the payload exactly how the serverless API wants it
+        payload = {
+            "inputs": f"<s>[INST] {user_prompt} [/INST]",
+            "parameters": {
+                "max_new_tokens": 250,
+                "temperature": 0.7,
+                "return_full_text": False
+            }
+        }
         
-        # Clean up output text by stripping out the original prompt block if echoed
-        reply_text = decoded.split(user_prompt)[-1].strip() if user_prompt in decoded else decoded
+        response = requests.post(API_URL, headers=headers, json=payload)
+        res_data = response.json()
         
-        return jsonify({"reply": reply_text})
+        # Check if the API tells us the model is still loading weights on their servers
+        if isinstance(res_data, dict) and "error" in res_data:
+            err_msg = res_data.get("error", "")
+            if "loading" in err_msg.lower():
+                est_time = round(res_data.get('estimated_time', 20))
+                return jsonify({"reply": f"SYS_BOOT: Krypton-1 is loading on Hugging Face clusters. Activation wait time: {est_time}s. Try hitting RUN again in a moment!"})
+            return jsonify({"reply": f"SYS_ALERT: API Error. Details: {err_msg}"})
+            
+        # Extract the text cleanly
+        if isinstance(res_data, list) and len(res_data) > 0:
+            reply = res_data[0].get("generated_text", "Data packet empty.")
+            return jsonify({"reply": reply})
+        else:
+            return jsonify({"reply": f"SYS_ALERT: Unexpected API data layout format. Received: {str(res_data)}"})
         
     except Exception as e:
-        return jsonify({"reply": f"SYS_ALERT: Internal processing breakdown. Details: {str(e)}"})
+        return jsonify({"reply": f"SYS_ERR: Proxy pipeline broken. {str(e)}"})
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)