| """Keyword / vocabulary extraction.""" |
| from __future__ import annotations |
|
|
| from typing import Tuple, List, Dict, Any |
|
|
| from src.config import KEYWORD_SYSTEM_PROMPT |
| from src.model_loader import get_model_and_tokenizer |
| from src.generation import chat_generate |
|
|
|
|
| def extract_keywords(text: str) -> Tuple[List[Dict[str, Any]], str]: |
| """Extract important vocabulary from Japanese text.""" |
| if not text or not text.strip(): |
| return [], "" |
|
|
| model, tokenizer = get_model_and_tokenizer() |
| messages = [ |
| {"role": "system", "content": KEYWORD_SYSTEM_PROMPT}, |
| {"role": "user", "content": text.strip()}, |
| ] |
| raw = chat_generate(model, tokenizer, messages, max_new_tokens=768, do_sample=False) |
| items = [] |
| for line in raw.strip().splitlines(): |
| line = line.strip() |
| if not line: |
| continue |
| if " — " in line: |
| head, meaning = line.split(" — ", 1) |
| if head.endswith(")") and " (" in head: |
| word, reading = head.rsplit(" (", 1) |
| reading = reading[:-1] |
| items.append({"word": word.strip(), "reading": reading.strip(), "meaning": meaning.strip()}) |
| else: |
| items.append({"word": head.strip(), "reading": "", "meaning": meaning.strip()}) |
| else: |
| items.append({"word": line, "reading": "", "meaning": ""}) |
| return items, "" |
|
|
