Update app.py

app.py

@@ -1,103 +1,122 @@
 import os
+import threading
 import requests
-from flask import Flask, request, jsonify
-from llama_cpp import Llama
 import subprocess
-import time
 import json
+from flask import Flask, jsonify, request
+from llama_cpp import Llama
 
 app = Flask(__name__)
 
-# Use /tmp directory for storing the model
-MODEL_DIR = "/tmp/model"
-MODEL_PATH = os.path.join(MODEL_DIR, "calme-3.3-llamaloi-3b.Q4_K_M.gguf")
-GH_PAT = "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"  # GitHub Personal Access Token
+# Configuration
+MODEL_DIR = "/data/model"  # Persistent storage directory
+MODEL_NAME = "calme-3.3-llamaloi-3b.Q4_K_M.gguf"
+MODEL_PATH = os.path.join(MODEL_DIR, MODEL_NAME)
+MODEL_URL = "https://huggingface.co/MaziyarPanahi/calme-3.3-llamaloi-3b-GGUF/resolve/main/calme-3.3-llamaloi-3b.Q4_K_M.gguf"
+GH_PAT = os.getenv("GH_PAT")
 REPO_URL = "https://github.com/NitinBot001/Audio-url-new-js.git"
 
-def download_model():
-    os.makedirs(MODEL_DIR, exist_ok=True)
-    if not os.path.exists(MODEL_PATH):
-        print("Downloading model...")
-        r = requests.get(
-            "https://huggingface.co/MaziyarPanahi/calme-3.3-llamaloi-3b-GGUF/resolve/main/calme-3.3-llamaloi-3b.Q4_K_M.gguf",
-            stream=True,
+# Global state
+initialized = False
+llm = None
+
+def background_init():
+    global initialized, llm
+    try:
+        # 1. Ensure model directory exists
+        os.makedirs(MODEL_DIR, exist_ok=True)
+        
+        # 2. Download model if not exists
+        if not os.path.exists(MODEL_PATH):
+            print("Downloading model...")
+            with requests.get(MODEL_URL, stream=True) as r:
+                r.raise_for_status()
+                with open(MODEL_PATH, "wb") as f:
+                    for chunk in r.iter_content(chunk_size=8192):
+                        f.write(chunk)
+            print("Model download complete")
+
+        # 3. Initialize LLM
+        llm = Llama(
+            model_path=MODEL_PATH,
+            n_ctx=8192,
+            n_threads=2,
+            n_gpu_layers=0,
+            verbose=False
         )
-        with open(MODEL_PATH, "wb") as f:
-            for chunk in r.iter_content(chunk_size=8192):
-                f.write(chunk)
+
+        # 4. Start tunnel and update repo
+        tunnel_url = start_tunnel()
+        update_repo_with_tunnel(tunnel_url)
+        
+        initialized = True
+        print("Initialization complete")
+    except Exception as e:
+        print(f"Initialization failed: {str(e)}")
 
 def start_tunnel():
-    tunnel_process = subprocess.Popen(
-        ["npx", "nport", "-s", "ai-service", "-p", "5000"],
+    """Start nport tunnel and return URL"""
+    proc = subprocess.Popen(
+        ["npx", "nport", "-s", "hf-space", "-p", "7860"],
         stdout=subprocess.PIPE,
         stderr=subprocess.PIPE,
+        text=True
     )
-    time.sleep(10)
-
-    tunnel_url = None
-    for line in iter(tunnel_process.stdout.readline, b""):
-        line = line.decode("utf-8").strip()
+    
+    # Wait for tunnel URL
+    while True:
+        line = proc.stdout.readline()
         if "your domain is:" in line:
-            tunnel_url = line.split("your domain is: ")[1]
+            return line.split("your domain is: ")[1].strip()
+        if proc.poll() is not None:
             break
+    raise RuntimeError("Failed to establish tunnel")
 
-    if not tunnel_url:
-        raise Exception("Failed to extract tunnel URL")
-    return tunnel_url
-
-def push_tunnel_url_to_repo(tunnel_url):
-    instance_data = {"tunnel_url": tunnel_url}
-    with open("/tmp/instance.json", "w") as f:
-        json.dump(instance_data, f)
-
-    repo_dir = "/tmp/repo"
-    repo_url = f"https://x-access-token:{GH_PAT}@github.com/NitinBot001/Audio-url-new-js.git"
+def update_repo_with_tunnel(url):
+    """Update GitHub repository with tunnel URL"""
+    repo_dir = "/data/repo"
+    instance_path = os.path.join(repo_dir, "instance.json")
     
+    # Clone or update repository
     if os.path.exists(repo_dir):
-        subprocess.run(["rm", "-rf", repo_dir], check=True)
-    
-    subprocess.run(["git", "clone", repo_url, repo_dir], check=True)
-    os.chdir(repo_dir)
+        subprocess.run(["git", "-C", repo_dir, "pull"], check=True)
+    else:
+        subprocess.run([
+            "git", "clone",
+            f"https://x-access-token:{GH_PAT}@github.com/NitinBot001/Audio-url-new-js.git",
+            repo_dir
+        ], check=True)
 
-    subprocess.run(["mv", "/tmp/instance.json", "."], check=True)
-    subprocess.run(["git", "config", "user.email", "[email protected]"], check=True)
-    subprocess.run(["git", "config", "user.name", "github-actions"], check=True)
-    subprocess.run(["git", "add", "instance.json"], check=True)
+    # Update instance.json
+    with open(instance_path, "w") as f:
+        json.dump({"tunnel_url": url}, f)
 
-    result = subprocess.run(["git", "status", "--porcelain"], stdout=subprocess.PIPE, text=True)
-    if result.stdout.strip():
-        subprocess.run(["git", "commit", "-m", f"Update tunnel URL to {tunnel_url}"], check=True)
-        subprocess.run(["git", "push", "origin", "main"], check=True)
-    else:
-        print("No changes to commit.")
+    # Commit and push changes
+    subprocess.run(["git", "-C", repo_dir, "add", "."], check=True)
+    subprocess.run([
+        "git", "-C", repo_dir,
+        "commit", "-m", f"Update tunnel URL: {url}"
+    ], check=True)
+    subprocess.run(["git", "-C", repo_dir, "push"], check=True)
 
 @app.route("/chat", methods=["GET"])
-def chat():
+def chat_endpoint():
+    if not initialized:
+        return jsonify({"error": "Service initializing"}), 503
     message = request.args.get("message", "")
-    prompt = (
-        f"<|begin_of_text|>"
-        f"<|start_header_id|>user<|end_header_id|>\n"
-        f"{message}"
-        f"<|eot_id|>\n"
-        f"<|start_header_id|>assistant<|end_header_id|>\n"
-    )
-    output = llm(
-        prompt,
-        max_tokens=2048,
-        stop=["<|eot_id|>"],
-        temperature=0.8,
-        top_p=0.9,
-    )
-    return jsonify({"response": output["choices"][0]["text"].strip()})
+    prompt = f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n{message}<|eot_id|>\n<|start_header_id|>assistant<|end_header_id|>\n"
+    output = llm(prompt, max_tokens=512, stop=["<|eot_id|>"])
+    return jsonify({"response": output['choices'][0]['text'].strip()})
+
+@app.route("/health")
+def health_check():
+    return jsonify({
+        "status": "ready" if initialized else "initializing",
+        "model_loaded": os.path.exists(MODEL_PATH)
+    }), 200 if initialized else 503
 
 if __name__ == "__main__":
-    download_model()
-    llm = Llama(
-        model_path=MODEL_PATH,
-        n_ctx=131072,
-        n_threads=2,
-        n_gpu_layers=0,
-        verbose=False,
-    )
-    tunnel_url = start_tunnel()
-    push_tunnel_url_to_repo(tunnel_url)
+    # Start initialization in background
+    threading.Thread(target=background_init, daemon=True).start()
+    # Start Flask server
+    app.run(host="0.0.0.0", port=7860)