File size: 2,705 Bytes
c7256ee | 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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | import os
import re
from huggingface_hub import InferenceClient
# the functions for resolving and generating titles
# it tries to query and hf model for title
# some shitty fallback logic, if models fail
# could improve candidate model defining code
def title_from_query(query: str) -> str:
stop_words = {
"a", "an", "and", "are", "as", "at", "be", "by", "can", "do", "for", "from", "how",
"i", "in", "is", "it", "me", "my", "of", "on", "or", "please", "show", "tell", "that",
"the", "this", "to", "we", "what", "when", "where", "which", "why", "with", "you", "your",
}
words = re.findall(r"[A-Za-z0-9][A-Za-z0-9\-_/+]*", query)
if not words:
return "New Chat"
filtered: list[str] = []
for word in words:
cleaned = word.strip("-_/+")
if not cleaned:
continue
if cleaned.lower() in stop_words:
continue
filtered.append(cleaned)
if len(filtered) >= 6:
break
chosen = filtered if filtered else words[:6]
normalized = [w.capitalize() if w.islower() else w for w in chosen]
title = " ".join(normalized).strip()
return title[:80] if title else "New Chat"
def clean_title_text(raw: str) -> str:
text = (raw or "").strip()
text = text.replace("\n", " ").replace("\r", " ")
text = re.sub(r"^[\"'`\s]+|[\"'`\s]+$", "", text)
text = re.sub(r"\s+", " ", text).strip()
words = text.split()
if len(words) > 8:
text = " ".join(words[:8])
return text[:80]
def title_from_hf(query: str, client: InferenceClient, model_id: str) -> str | None:
system_prompt = (
"You generate short chat titles. Return only a title, no punctuation at the end, no quotes."
)
user_prompt = (
"Create a concise 3-7 word title for this user request:\n"
f"{query}"
)
response = client.chat_completion(
model=model_id,
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
],
max_tokens=24,
temperature=0.3,
)
if not response or not response.choices:
return None
raw_title = response.choices[0].message.content or ""
title = clean_title_text(raw_title)
if not title or title.lower() == "new chat":
return None
return title
def parse_title_model_candidates() -> list[str]:
raw = os.getenv(
"TITLE_MODEL_IDS",
"Qwen/Qwen2.5-1.5B-Instruct,CohereLabs/tiny-aya-global,meta-llama/Meta-Llama-3-8B-Instruct",
)
models = [m.strip() for m in raw.split(",") if m.strip()]
return models or ["meta-llama/Meta-Llama-3-8B-Instruct"]
|