Update app/main.py

app/main.py

@@ -1,75 +1,55 @@
+from fastapi import FastAPI, Request
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+from peft import PeftModel
+import torch
 import os
 import gdown
-import re
-import torch
-from fastapi import FastAPI, Request
-from pydantic import BaseModel
-from peft import PeftModel, PeftConfig
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
 
 app = FastAPI()
 
-DRIVE_FOLDER_URL = "https://drive.google.com/drive/folders/1S9xT92Zm9rZ4RSCxAe_DLld8vu78mqW4"
-LOCAL_ADAPTER_DIR = "adapter"
-BASE_MODEL = "Qwen/Qwen2-0.5B-Instruct"
-
-class PromptRequest(BaseModel):
-    prompt: str
+# Auto-download adapter from Google Drive (if not already present)
+ADAPTER_DIR = "adapter"
+ADAPTER_PATH = os.path.join(ADAPTER_DIR, "adapter_model.safetensors")
+DRIVE_FILE_ID = "1wnuE5t_m4ojI7YqxXZ8lBdtDFoHJJ6_H"  # version 1 model
 
-def download_latest_adapter():
-    print("🔽 Downloading adapter folder from Google Drive...")
-    gdown.download_folder(url=DRIVE_FOLDER_URL, output="gdrive_tmp", quiet=False, use_cookies=False)
+if not os.path.exists(ADAPTER_PATH):
+    os.makedirs(ADAPTER_DIR, exist_ok=True)
+    gdown.download(f"https://drive.google.com/uc?id={DRIVE_FILE_ID}", ADAPTER_PATH, quiet=False)
 
-    all_versions = sorted(
-        [d for d in os.listdir("gdrive_tmp") if re.match(r"version \d+", d)],
-        key=lambda x: int(x.split()[-1])
-    )
-    if not all_versions:
-        raise ValueError("❌ No version folders found in Google Drive folder.")
+# Load base model
+base_model = AutoModelForCausalLM.from_pretrained(
+    "Qwen/Qwen2-0.5B-Instruct",
+    device_map="auto",
+    torch_dtype=torch.float16
+)
 
-    latest = all_versions[-1]
-    src = os.path.join("gdrive_tmp", latest)
-    print(f"✅ Latest adapter found: {latest}")
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
 
-    os.makedirs(LOCAL_ADAPTER_DIR, exist_ok=True)
-    for file in os.listdir(src):
-        src_file = os.path.join(src, file)
-        dest_file = os.path.join(LOCAL_ADAPTER_DIR, file)
-        os.system(f"cp '{src_file}' '{dest_file}'")
+# Load LoRA adapter
+model = PeftModel.from_pretrained(base_model, ADAPTER_DIR)
+model.eval()
 
-    print(f"✅ Adapter copied to: {LOCAL_ADAPTER_DIR}")
+@app.post("/chat")
+async def chat(request: Request):
+    data = await request.json()
+    prompt = data.get("prompt")
 
-def load_model():
-    print("🚀 Loading base model...")
-    model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, device_map="auto", torch_dtype=torch.float16)
-    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
-    
-    print("🔗 Loading adapter...")
-    model = PeftModel.from_pretrained(model, LOCAL_ADAPTER_DIR)
-    model.eval()
+    if not prompt:
+        return {"error": "No prompt provided."}
 
-    return model, tokenizer
-
-# Step 1: Download latest adapter
-download_latest_adapter()
-
-# Step 2: Load model and tokenizer
-model, tokenizer = load_model()
-
-@app.post("/generate")
-async def generate_text(request: PromptRequest):
-    prompt = request.prompt.strip()
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.cuda()
+    full_prompt = f"<|im_start|>system\nYou are a helpful assistant<|im_end|>\n<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
 
+    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
     with torch.no_grad():
         outputs = model.generate(
-            input_ids,
-            max_new_tokens=300,
-            do_sample=True,
+            **inputs,
+            max_new_tokens=256,
             temperature=0.7,
-            top_p=0.95,
-            eos_token_id=tokenizer.eos_token_id,
+            do_sample=True,
+            top_p=0.9
         )
 
-    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return {"response": result[len(prompt):].strip()}
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    response = response.split("<|im_start|>assistant\n")[-1].strip()
+    return {"response": response}