Update app.py

app.py

@@ -1,34 +1,35 @@
 #!/usr/bin/env python3
-# app.py — Full Restoration-to-ChatGPT-Clone for Hugging Face
-# ⚙️ Zero mocks · Zero OAuth · Handles 500 MB+ OpenAI export ZIP
+# app.py — Full OpenAI JSON Restoration + Persistent Chat Replica
+# 🧠 Real memory · Real UI · No mocks · No placeholders
 
-import os, zipfile, json, tempfile, gc, time
+import os, zipfile, json, tempfile, time, gc, shelve
 from pathlib import Path
 import gradio as gr
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import torch
 
-# === CONFIGURATION ===
+# ================= CONFIG =================
 MODEL_ID = os.environ.get("HF_MODEL", "mistralai/Mistral-7B-Instruct-v0.2")
-DEVICE = "cuda" if os.environ.get("USE_CUDA", "0") == "1" else "cpu"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+MEMORY_PATH = "chat_memory.db"
 
-# === LOAD MODEL ===
-print(f"⏳ Loading model {MODEL_ID} on {DEVICE} ...")
+# ================= LOAD MODEL =================
+print(f"⏳ Loading {MODEL_ID} on {DEVICE}")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto", device_map="auto")
 llm = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto")
-print("✅ Model loaded")
+print("✅ Model ready")
 
-# === DATA PIPELINE ===
+# ================= DATA PIPELINE =================
 def unpack_zip(zip_file):
-    """Unpack the user-uploaded OpenAI ZIP to a temp dir and return file list"""
-    tmp_dir = tempfile.mkdtemp(prefix="chat_restore_")
+    """Unpack uploaded OpenAI ZIP safely."""
+    tmp_dir = tempfile.mkdtemp(prefix="restore_")
     with zipfile.ZipFile(zip_file.name, "r") as zf:
         zf.extractall(tmp_dir)
-    paths = list(Path(tmp_dir).rglob("*.json"))
-    return tmp_dir, paths
+    return list(Path(tmp_dir).rglob("*.json")), tmp_dir
 
 def restore_conversations(paths):
-    """Stream large OpenAI export JSONs into a chronological chat list"""
+    """Parse and chronologically order OpenAI JSON messages."""
     chats, count = [], 0
     for p in paths:
         try:
@@ -39,7 +40,7 @@ def restore_conversations(paths):
                     msg = v.get("message")
                     if not msg or not msg.get("author"): 
                         continue
-                    role = msg["author"].get("role")
+                    role = msg["author"].get("role", "")
                     parts = msg.get("content", {}).get("parts")
                     text = parts[0] if parts else ""
                     if text.strip():
@@ -47,56 +48,75 @@ def restore_conversations(paths):
                         count += 1
         except Exception:
             continue
-    chats.sort(key=lambda x: x["role"] != "system")  # system first
     gc.collect()
     return chats, count
 
-def summarize_conversations(chats, limit=50):
-    """Create a condensed snapshot for quick inspection"""
-    return [{"role": c["role"], "text": c["text"][:180]} for c in chats[:limit]]
-
-# === CHAT ENGINE ===
-def reply(history, user_msg):
-    prompt = ""
-    for speaker, msg in history[-8:]:
-        prompt += f"{speaker}: {msg}\n"
-    prompt += f"user: {user_msg}\nassistant:"
-    out = llm(prompt, max_new_tokens=180, do_sample=True, temperature=0.7, top_p=0.9)[0]["generated_text"]
+def summarize(chats, limit=50):
+    """Condensed preview."""
+    return [{"role": c["role"], "text": c["text"][:200]} for c in chats[:limit]]
+
+# ================= MEMORY =================
+def load_memory():
+    with shelve.open(MEMORY_PATH) as db:
+        return db.get("history", [])
+def save_memory(history):
+    with shelve.open(MEMORY_PATH) as db:
+        db["history"] = history
+
+# ================= CHAT ENGINE =================
+def chat_reply(history, user_input):
+    full_prompt = "\n".join([f"{r[0]}: {r[1]}" for r in history[-12:]]) + f"\nuser: {user_input}\nassistant:"
+    out = llm(full_prompt, max_new_tokens=200, temperature=0.7, top_p=0.9)[0]["generated_text"]
     answer = out.split("assistant:")[-1].strip()
-    history.append(("user", user_msg))
+    history.append(("user", user_input))
     history.append(("assistant", answer))
+    save_memory(history)
     return history, ""
 
-# === UI BUILD ===
-theme = gr.themes.Soft(primary_hue="orange", neutral_hue="black")
+# ================= UI THEME =================
+neomorph = gr.themes.Base(
+    primary_hue="orange",
+    secondary_hue="black",
+    neutral_hue="black",
+    font="IBM Plex Sans",
+    radius_size=gr.themes.sizes.radius_sm,
+    shadow_drop="drop-lg",
+)
 
-with gr.Blocks(title="ChatGPT Restoration Engine", theme=theme) as app:
-    gr.Markdown("# 🧱 ChatGPT Restoration Engine")
-    gr.Markdown("Upload your exported OpenAI JSON ZIP to rebuild a full offline chat environment.")
+# ================= GRADIO UI =================
+with gr.Blocks(theme=neomorph, title="ChatGPT Memory Restoration") as app:
+    gr.HTML(
+        "<h1 style='text-align:center;color:#FFA500;'>💬 ChatGPT Memory Restoration</h1>"
+        "<p style='text-align:center;color:#888;'>Upload your OpenAI export → restore → chat with persistent memory</p>"
+    )
 
     with gr.Tab("📦 Restore Export"):
-        zip_in = gr.File(label="Upload OpenAI ZIP (up to 500 MB)", file_types=[".zip"])
-        summary_out = gr.JSON(label="Summary / Preview (Top 50 messages)")
-        restore_btn = gr.Button("Restore Conversations", variant="primary")
+        zip_in = gr.File(label="Upload your OpenAI Export ZIP (≤500 MB)", file_types=[".zip"])
+        preview_out = gr.JSON(label="Preview (Top 50 messages)")
+        restore_btn = gr.Button("🧱 Rebuild Memory", variant="primary")
 
-    with gr.Tab("💬 Recreated Chat"):
-        chatbot = gr.Chatbot(label="Your Restored ChatGPT Clone")
-        user_input = gr.Textbox(label="Message")
-        send_btn = gr.Button("Send")
-
-    state = gr.State([])
+    with gr.Tab("🧠 Chat Interface"):
+        chatbot = gr.Chatbot(label="Persistent Chat", show_label=False, height=550)
+        msg_box = gr.Textbox(label="Message", placeholder="Type something…", autofocus=True)
+        send_btn = gr.Button("Send", variant="primary")
+        clear_btn = gr.Button("Clear Memory", variant="secondary")
+        history_state = gr.State(load_memory())
 
+    # ===== RESTORE HANDLER =====
     def handle_restore(zip_in):
-        t0 = time.time()
-        tmp_dir, paths = unpack_zip(zip_in)
+        start = time.time()
+        paths, tmp_dir = unpack_zip(zip_in)
         chats, count = restore_conversations(paths)
-        snapshot = summarize_conversations(chats)
-        print(f"🧾 Restored {count} messages from {len(paths)} files in {time.time()-t0:.1f}s")
-        return snapshot, chats
+        save_memory([(c["role"], c["text"]) for c in chats])
+        print(f"Restored {count} messages from {len(paths)} files in {time.time()-start:.1f}s")
+        return summarize(chats), [(c["role"], c["text"]) for c in chats]
+
+    restore_btn.click(fn=handle_restore, inputs=[zip_in], outputs=[preview_out, history_state])
 
-    restore_btn.click(fn=handle_restore, inputs=[zip_in], outputs=[summary_out, state])
-    send_btn.click(fn=reply, inputs=[state, user_input], outputs=[state, user_input])
+    # ===== CHAT HANDLERS =====
+    send_btn.click(fn=chat_reply, inputs=[history_state, msg_box], outputs=[history_state, msg_box])
+    clear_btn.click(lambda: ([], ""), None, [history_state, msg_box])
 
-# === DEPLOY ===
+# ================= RUN =================
 if __name__ == "__main__":
     app.launch(server_name="0.0.0.0", server_port=7860, share=True)