Julian Carlson
67 lines
2.4 KiB
Python
67 lines
2.4 KiB
Python
"""Extract topic labels for chunks using a local LLM."""
|
|
import requests
|
|
import os
|
|
from .config import CONFIG
|
|
|
|
# Use a fast model — chosen at config time
|
|
TOPIC_MODEL = os.getenv("KB_TOPIC_MODEL", "qwen2.5-coder:7b")
|
|
|
|
|
|
def extract_topic(chunk_text, doc_title=None, doc_category=None):
|
|
"""Generate a brief topic label for a chunk.
|
|
|
|
Returns a single-line description of what specific entities,
|
|
concepts, and topics this chunk discusses.
|
|
"""
|
|
# Truncate very long chunks for the LLM
|
|
if len(chunk_text) > 3000:
|
|
chunk_text = chunk_text[:3000] + "..."
|
|
|
|
context_hint = ""
|
|
if doc_title:
|
|
context_hint += f"\nThe broader document is titled: {doc_title}"
|
|
if doc_category:
|
|
context_hint += f"\nThe document category: {doc_category}"
|
|
|
|
prompt = (
|
|
f"<|im_start|>system\n"
|
|
f"You extract topic labels for retrieval. Given a passage, identify the specific "
|
|
f"entities, concepts, names, and topics it discusses. Output a single concise line "
|
|
f"of 10-30 words listing what this passage is specifically about. Include proper "
|
|
f"nouns, technical terms, and metric values. Do not summarize the content; just "
|
|
f"label what it covers.<|im_end|>\n"
|
|
f"<|im_start|>user\n"
|
|
f"Passage:{context_hint}\n\n{chunk_text}<|im_end|>\n"
|
|
f"<|im_start|>assistant\n"
|
|
)
|
|
|
|
try:
|
|
resp = requests.post(
|
|
f"{CONFIG.ollama_url}/api/generate",
|
|
json={
|
|
"model": TOPIC_MODEL,
|
|
"prompt": prompt,
|
|
"stream": False,
|
|
"options": {
|
|
"num_predict": 60,
|
|
"temperature": 0.0,
|
|
"top_p": 1.0,
|
|
},
|
|
"raw": True,
|
|
},
|
|
timeout=60,
|
|
)
|
|
resp.raise_for_status()
|
|
text = resp.json().get("response", "").strip()
|
|
|
|
# Take only the first line, drop any "Labels:" prefix
|
|
first_line = text.split("\n")[0].strip()
|
|
for prefix in ("Topics:", "Topic:", "Labels:", "Label:", "Subject:"):
|
|
if first_line.lower().startswith(prefix.lower()):
|
|
first_line = first_line[len(prefix):].strip()
|
|
|
|
return first_line[:300] # cap length
|
|
except Exception as e:
|
|
print(f"Topic extraction error: {e}", flush=True)
|
|
return ""
|