Update app/main.py

app/main.py

@@ -1,89 +1,62 @@
-# app/main.py
-import os
-import torch
-import gdown
-import re
-import shutil
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from peft import PeftModel, PeftConfig
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
+import os
 
-# ====== CONFIG ======
-DRIVE_FOLDER_URL = "https://drive.google.com/drive/folders/1S9xT92Zm9rZ4RSCxAe_DLld8vu78mqW4"
-ADAPTER_DIR = "adapter"
-TEMP_DIR = "gdrive_tmp"
-BASE_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
+# === CONFIG ===
+HF_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
+ADAPTER_PATH = "adapter"  # folder where your LoRA is saved
+API_KEY = os.getenv("API_KEY", "your-secret-key")  # Set in HF Space secrets
 
-# ====== FASTAPI SETUP ======
+# === FastAPI Setup ===
 app = FastAPI()
 
-class Message(BaseModel):
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # adjust if needed
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# === Load Model & Tokenizer (CPU only) ===
+print("🔧 Loading model on CPU...")
+tokenizer = AutoTokenizer.from_pretrained(HF_MODEL, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(HF_MODEL, torch_dtype=torch.float32, trust_remote_code=True)
+model = PeftModel.from_pretrained(model, ADAPTER_PATH)
+model = model.to("cpu")
+model.eval()
+print("✅ Model ready on CPU.")
+
+# === Request Schema ===
+class ChatRequest(BaseModel):
     prompt: str
+    api_key: str
 
-# ====== DOWNLOAD LATEST ADAPTER ======
-def download_latest_adapter():
-    print("🔽 Downloading adapter folder from Google Drive...")
-    gdown.download_folder(url=DRIVE_FOLDER_URL, output=TEMP_DIR, quiet=False, use_cookies=False)
-
-    all_versions = sorted(
-        [d for d in os.listdir(TEMP_DIR) if re.match(r"version \d+", d)],
-        key=lambda x: int(x.split()[-1])
-    )
-    if not all_versions:
-        raise ValueError("❌ No adapter versions found.")
-
-    latest = all_versions[-1]
-    src = os.path.join(TEMP_DIR, latest)
-
-    os.makedirs(ADAPTER_DIR, exist_ok=True)
-    for f in os.listdir(ADAPTER_DIR):
-        os.remove(os.path.join(ADAPTER_DIR, f))
-
-    for f in os.listdir(src):
-        shutil.copy(os.path.join(src, f), os.path.join(ADAPTER_DIR, f))
-
-    print(f"✅ Adapter '{latest}' copied to '{ADAPTER_DIR}'")
-
-# ====== LOAD MODEL ======
-def load_model():
-    print("🔧 Loading base model...")
-    base_model = AutoModelForCausalLM.from_pretrained(
-        BASE_MODEL,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        device_map="auto"
-    )
-
-    print("🔗 Loading LoRA adapter...")
-    model = PeftModel.from_pretrained(base_model, ADAPTER_DIR)
-
-    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
-    return model, tokenizer
-
-# ====== RUN ======
-download_latest_adapter()
-model, tokenizer = load_model()
+@app.get("/")
+def root():
+    return {"message": "✅ Qwen2.5 Chat API running."}
 
 @app.post("/chat")
-def chat(msg: Message):
-    prompt = msg.prompt.strip()
-
-    messages = [
-        {"role": "user", "content": prompt}
-    ]
-    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-
-    inputs = tokenizer(text, return_tensors="pt").to(model.device)
-    with torch.no_grad():
-        output = model.generate(
-            **inputs,
-            max_new_tokens=512,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.9
-        )
-
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-    response = response.replace(text, "").strip()
+def chat(req: ChatRequest):
+    if req.api_key != API_KEY:
+        raise HTTPException(status_code=401, detail="Invalid API Key")
+
+    input_text = f"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n{req.prompt}<|im_end|>\n<|im_start|>assistant\n"
+
+    inputs = tokenizer(input_text, return_tensors="pt").to("cpu")
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=512,
+        temperature=0.7,
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id
+    )
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    return {"response": response}
+    # Extract response after assistant tag
+    final_resp = response.split("<|im_start|>assistant\n")[-1].strip()
+    return {"response": final_resp}