Update app/main.py

app/main.py

@@ -1,55 +1,49 @@
-from fastapi import FastAPI, Request
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from peft import PeftModel
 import torch
-import os
 import gdown
+import os
+import zipfile
+
+# Constants
+BASE_MODEL = "Qwen/Qwen2-0.5B-Instruct"
+ADAPTER_FOLDER = "adapter"
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
 
-app = FastAPI()
+# Step 1: Download adapter zip from Drive (version 1)
+zip_url = "https://drive.google.com/uc?id=1z8U98kW9GD29t-3v8LDu0SsdqJ_vzNvQ"  # Your .zip file link
+zip_path = "adapter.zip"
 
-# 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
+if not os.path.exists(ADAPTER_FOLDER):
+    print("📥 Downloading adapter...")
+    gdown.download(zip_url, zip_path, quiet=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)
+    print("📂 Extracting adapter...")
+    with zipfile.ZipFile(zip_path, "r") as zip_ref:
+        zip_ref.extractall(ADAPTER_FOLDER)
 
-# Load base model
+# Step 2: Load base model (non-quantized, CPU-friendly)
+print("🚀 Loading base model...")
 base_model = AutoModelForCausalLM.from_pretrained(
-    "Qwen/Qwen2-0.5B-Instruct",
+    BASE_MODEL,
+    torch_dtype=torch.float16,
     device_map="auto",
-    torch_dtype=torch.float16
+    token=HF_TOKEN
 )
 
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
-
-# Load LoRA adapter
-model = PeftModel.from_pretrained(base_model, ADAPTER_DIR)
-model.eval()
-
-@app.post("/chat")
-async def chat(request: Request):
-    data = await request.json()
-    prompt = data.get("prompt")
-
-    if not prompt:
-        return {"error": "No prompt provided."}
+# Step 3: Apply LoRA adapter
+print("🔧 Applying LoRA adapter...")
+model = PeftModel.from_pretrained(base_model, ADAPTER_FOLDER)
 
-    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"
+# Step 4: Load tokenizer
+print("🧠 Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
 
-    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=256,
-            temperature=0.7,
-            do_sample=True,
-            top_p=0.9
-        )
+# Step 5: Inference pipeline
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    response = response.split("<|im_start|>assistant\n")[-1].strip()
-    return {"response": response}
+# Step 6: Try a prompt
+prompt = "What is the capital of India?"
+print("💬 Prompt:", prompt)
+output = pipe(prompt, max_new_tokens=100, do_sample=True, temperature=0.7)
+print("📤 Output:", output[0]["generated_text"])