Update app/main.py

app/main.py

@@ -1,57 +1,75 @@
+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
-from peft import PeftModel
-import json
-import os
 
-# Load tokenizer and base model
-base_model = "Qwen/Qwen2-0.5B-Instruct"
-tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    base_model,
-    device_map="cuda" if torch.cuda.is_available() else "cpu",
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    trust_remote_code=True
-)
-
-# Clean up adapter_config.json before loading adapter
-adapter_config_path = "./adapter/adapter_config.json"
-if os.path.exists(adapter_config_path):
-    with open(adapter_config_path, "r") as f:
-        adapter_config = json.load(f)
-    for key in ["corda_config", "eva_config", "megatron_config"]:
-        adapter_config.pop(key, None)
-    with open(adapter_config_path, "w") as f:
-        json.dump(adapter_config, f)
-
-# Load adapter
-model = PeftModel.from_pretrained(model, "./adapter", is_trainable=False)
-model.eval()
-
-# Simple chat function
-def chat(prompt):
-    messages = [
-        {"role": "system", "content": "You are a helpful assistant."},
-        {"role": "user", "content": prompt}
-    ]
-    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
-
-    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    generated_ids = model.generate(
-        **model_inputs,
-        max_new_tokens=512,
-        do_sample=True,
-        temperature=0.7,
-        streamer=streamer
+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
+
+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)
+
+    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])
     )
-    output = tokenizer.decode(generated_ids[0][model_inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
-    return output
-
-# Example
-if __name__ == "__main__":
-    while True:
-        prompt = input("User: ")
-        if prompt.lower() in ["exit", "quit"]:
-            break
-        print("AI:", chat(prompt))
+    if not all_versions:
+        raise ValueError("❌ No version folders found in Google Drive folder.")
+
+    latest = all_versions[-1]
+    src = os.path.join("gdrive_tmp", latest)
+    print(f"✅ Latest adapter found: {latest}")
+
+    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}'")
+
+    print(f"✅ Adapter copied to: {LOCAL_ADAPTER_DIR}")
+
+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()
+
+    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()
+
+    with torch.no_grad():
+        outputs = model.generate(
+            input_ids,
+            max_new_tokens=300,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.95,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+
+    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return {"response": result[len(prompt):].strip()}