Update app.py

app.py

@@ -1,15 +1,11 @@
 """
 Codette AI Space — FastAPI + streaming chat API
-Compatible with the Ollama /api/chat streaming format so the HTML widget
-needs only a URL change to work.
-
-Adapter files should live in ./adapter/ inside this Space repo.
-Base model: meta-llama/Llama-3.2-1B
 """
 
 import json
 import asyncio
 import threading
+import os
 from pathlib import Path
 from typing import Iterator
 
@@ -20,22 +16,14 @@ from fastapi.responses import StreamingResponse, HTMLResponse
 from peft import PeftModel
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
-# ── Config ───────────────────────────────────────────────────────────────────
 BASE_MODEL   = "meta-llama/Llama-3.2-1B"
 ADAPTER_PATH = Path(__file__).parent / "adapter"
 DEVICE       = "cuda" if torch.cuda.is_available() else "cpu"
+HF_TOKEN     = os.environ.get("HF_TOKEN")  # set as a Secret in Space settings
 
-# ── App ───────────────────────────────────────────────────────────────────────
 app = FastAPI(title="Codette AI")
+app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],   # Squarespace domain — keep open so the widget works
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# ── Model (loaded once at startup) ────────────────────────────────────────────
 print(f"Loading tokenizer from {ADAPTER_PATH} …")
 tokenizer = AutoTokenizer.from_pretrained(str(ADAPTER_PATH))
 if tokenizer.pad_token is None:
@@ -47,24 +35,21 @@ base = AutoModelForCausalLM.from_pretrained(
     torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
     device_map="auto",
     low_cpu_mem_usage=True,
+    token=HF_TOKEN,
 )
 
-print(f"Loading LoRA adapter from {ADAPTER_PATH} …")
+print(f"Loading LoRA adapter …")
 model = PeftModel.from_pretrained(base, str(ADAPTER_PATH))
-
-print("Merging LoRA weights into base model …")
-model = model.merge_and_unload()   # ← this is the actual merge step
+print("Merging LoRA weights …")
+model = model.merge_and_unload()
 model.eval()
 print(f"✅ Model ready on {DEVICE}")
 
 
-# ── Helpers ───────────────────────────────────────────────────────────────────
-def build_prompt(messages: list[dict]) -> str:
-    """Convert OpenAI-style messages to a simple Llama-3.2 instruct prompt."""
+def build_prompt(messages):
     parts = []
     for m in messages:
-        role = m.get("role", "user")
-        content = m.get("content", "")
+        role, content = m.get("role", "user"), m.get("content", "")
         if role == "system":
             parts.append(f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n{content}<|eot_id|>")
         elif role == "user":
@@ -75,33 +60,17 @@ def build_prompt(messages: list[dict]) -> str:
     return "".join(parts)
 
 
-def stream_tokens(messages: list[dict], max_new_tokens: int = 512) -> Iterator[str]:
-    prompt = build_prompt(messages)
-    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
-
-    streamer = TextIteratorStreamer(
-        tokenizer, skip_prompt=True, skip_special_tokens=True
-    )
-
-    gen_kwargs = dict(
-        **inputs,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        temperature=0.7,
-        top_p=0.9,
-        streamer=streamer,
-    )
-
-    thread = threading.Thread(target=model.generate, kwargs=gen_kwargs)
+def stream_tokens(messages, max_new_tokens=512):
+    inputs = tokenizer(build_prompt(messages), return_tensors="pt").to(DEVICE)
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    thread = threading.Thread(target=model.generate, kwargs=dict(**inputs, max_new_tokens=max_new_tokens,
+        do_sample=True, temperature=0.7, top_p=0.9, streamer=streamer))
     thread.start()
-
     for token in streamer:
         yield token
-
     thread.join()
 
 
-# ── Routes ────────────────────────────────────────────────────────────────────
 @app.get("/", response_class=HTMLResponse)
 async def root():
     return "<h2>Codette AI is running ✅</h2><p>POST /api/chat to chat.</p>"
@@ -111,36 +80,11 @@ async def root():
 async def chat(request: Request):
     body = await request.json()
     messages = body.get("messages", [])
-    stream = body.get("stream", True)
 
-    if not stream:
-        # Non-streaming — collect everything first
-        full = "".join(stream_tokens(messages))
-        return {
-            "message": {"role": "assistant", "content": full},
-            "done": True,
-        }
-
-    # Streaming — mimic Ollama's NDJSON format exactly
     async def event_stream():
-        full = ""
         for token in stream_tokens(messages):
-            full += token
-            chunk = json.dumps({
-                "message": {"role": "assistant", "content": token},
-                "done": False,
-            })
-            yield chunk + "\n"
-            await asyncio.sleep(0)   # yield control to event loop
-
-        # Final message with done=true
-        yield json.dumps({
-            "message": {"role": "assistant", "content": ""},
-            "done": True,
-        }) + "\n"
-
-    return StreamingResponse(
-        event_stream(),
-        media_type="application/x-ndjson",
-        headers={"X-Accel-Buffering": "no"},
-    )
+            yield json.dumps({"message": {"role": "assistant", "content": token}, "done": False}) + "\n"
+            await asyncio.sleep(0)
+        yield json.dumps({"message": {"role": "assistant", "content": ""}, "done": True}) + "\n"
+
+    return StreamingResponse(event_stream(), media_type="application/x-ndjson", headers={"X-Accel-Buffering": "no"})