|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
download_latest_adapter() |
|
|
|
|
|
|
|
|
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) |
|
|
|
|
|
|
|
|
app = FastAPI() |
|
|
|
|
|
app.add_middleware( |
|
|
CORSMiddleware, |
|
|
allow_origins=["*"], |
|
|
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> |
|
|
""" |
|
|
|
