Spaces:

Ais203
/

aigen

Sleeping

App Files Files Community

Ais commited on 26 days ago

Commit

90ddcea

verified ·

1 Parent(s): 363fc6a

Update app/main.py

Browse files

Files changed (1) hide show

app/main.py +78 -69

app/main.py CHANGED Viewed

@@ -1,80 +1,89 @@
 # app/main.py
-from fastapi import FastAPI, Form
-from fastapi.responses import HTMLResponse
-from fastapi.middleware.cors import CORSMiddleware
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from peft import PeftModel
-import torch
 import os
-from app.download_adapter import download_latest_adapter
-# === Step 1: Download Adapter ===
-download_latest_adapter()
-# === Step 2: Load Model and Tokenizer ===
-BASE_MODEL = "Qwen/Qwen2-0.5B-Instruct"
-ADAPTER_FOLDER = "adapter"
-HF_TOKEN = os.environ.get("HF_TOKEN", None)
-print("🚀 Loading base model...")
-base_model = AutoModelForCausalLM.from_pretrained(
-    BASE_MODEL,
-    torch_dtype=torch.float16,
-    device_map="auto",
-    token=HF_TOKEN,
-    trust_remote_code=True
-)
-print("🔧 Applying LoRA adapter...")
-model = PeftModel.from_pretrained(base_model, ADAPTER_FOLDER)
-print("🧠 Loading tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
-# === Step 3: FastAPI App ===
-app = FastAPI()
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],  # Allow all origins for testing
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-@app.get("/", response_class=HTMLResponse)
-async def form():
-    return """
-    <html>
-        <head><title>Qwen Chat</title></head>
-        <body>
-            <h2>Ask something:</h2>
-            <form method="post">
-                <textarea name="prompt" rows="4" cols="60"></textarea><br>
-                <input type="submit" value="Generate">
-            </form>
-        </body>
-    </html>
-    """
-@app.post("/", response_class=HTMLResponse)
-async def generate(prompt: str = Form(...)):
-    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"
-    output = pipe(full_prompt, max_new_tokens=256, do_sample=True, temperature=0.7)
-    response = output[0]["generated_text"].split("<|im_start|>assistant\n")[-1].strip()
-    return f"""
-    <html>
-        <head><title>Qwen Chat</title></head>
-        <body>
-            <h2>Your Prompt:</h2>
-            <p>{prompt}</p>
-            <h2>Response:</h2>
-            <p>{response}</p>
-            <a href="/">Ask again</a>
-        </body>
-    </html>
-    """

 # app/main.py
 import os
+import torch
+import gdown
+import re
+import shutil
+from fastapi import FastAPI, Request
+from pydantic import BaseModel
+from peft import PeftModel, PeftConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
+# ====== 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"
+# ====== FASTAPI SETUP ======
+app = FastAPI()
+class Message(BaseModel):
+    prompt: 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.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()
+    return {"response": response}