File size: 2,097 Bytes
474fd08 b62ffbd 474fd08 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import gradio as gr
# -------------------------
# Base + Adapter configuration
# -------------------------
base_model_name = "Qwen/Qwen3-4B-Instruct-2507"
adapter_model_name = "help2opensource/Qwen3-4B-Instruct-2507_mental_health_therapy"
device = "cuda" if torch.cuda.is_available() else "cpu"
# -------------------------
# Load base model and tokenizer
# -------------------------
tokenizer = AutoTokenizer.from_pretrained(base_model_name)
base_model = AutoModelForCausalLM.from_pretrained(
base_model_name,
torch_dtype=torch.float16 if device == "cuda" else torch.float32,
).to(device)
# -------------------------
# Load LoRA adapter
# -------------------------
model = PeftModel.from_pretrained(base_model, adapter_model_name)
# Optional: merge LoRA weights for faster inference
model = model.merge_and_unload()
def predict(message, history):
# Ensure history format is consistent
messages = history + [{"role": "user", "content": message}]
# Apply chat template correctly
try:
input_text = tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
except TypeError:
# For older tokenizers that don't support add_generation_prompt
input_text = tokenizer.apply_chat_template(messages, tokenize=False)
inputs = tokenizer(input_text, return_tensors="pt").to(device)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=1024,
temperature=0.7,
top_p=0.9,
do_sample=True,
)
decoded = tokenizer.decode(outputs[0], skip_special_tokens=False)
# Extract only the assistant’s final response
if "<|im_start|>assistant" in decoded:
response = (
decoded.split("<|im_start|>assistant")[-1].split("<|im_end|>")[0].strip()
)
else:
response = decoded
return response
demo = gr.ChatInterface(predict, type="messages")
demo.launch()
|