moodmate-api / app /emotion_detector.py
nyarderr's picture
Update app/emotion_detector.py
3bbe154 verified
import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
MODEL_NAME = "Qwen/Qwen1.5-0.5B"
LORA_REPO = "nyarderr/qwen-lora-goemotions"
## load tokenizer
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen1.5-0.5B", trust_remote_code=True)
## load base model
base_model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen1.5-0.5B",
trust_remote_code=True,
torch_dtype=torch.float16,
device_map="auto",
)
## load peft model
model = PeftModel.from_pretrained(
base_model,
LORA_REPO,
torch_dtype=torch.float16,
device_map="auto",
)
# Load LoRA adapter
model = PeftModel.from_pretrained(base_model, LORA_REPO)
model.eval()
def detect_emotion(text: str) -> str:
prompt = (
"Instrcution: Identify the emotion of the following text.\n"
f"Text: {text}\n"
"Emotion:"
)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
outputs = model.generate(**inputs, max_new_tokens=5, do_sample=False)
decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
emotion = decoded.split("Emotion:")[-1].strip().split()[0]
return emotion