Update app/deps.py

app/deps.py

@@ -4,10 +4,10 @@ import faiss, os, pickle, torch
 
 from app.config import settings
 
-# --- data ディレクトリの存在を保証（初回起動での書き込み失敗を防止） ---
+# data ディレクトリの存在を保証
 os.makedirs(settings.DATA_DIR, exist_ok=True)
 
-# Embeddings (multilingual-e5)
+# Embeddings
 emb_model = SentenceTransformer(settings.EMB_MODEL)
 
 # FAISS index (create or load)
@@ -20,9 +20,9 @@ if os.path.exists(index_path) and os.path.exists(store_path):
         store = pickle.load(f)
 else:
     index = faiss.IndexFlatIP(emb_model.get_sentence_embedding_dimension())
-    store = []  # list[dict]: {text, source_url, title, doc_id, chunk_id, score...}
+    store = []  # [{text, source_url, title, doc_id, chunk_id, ...}]
 
-# Generation model (Japanese-capable small model by default)
+# Text-generation model
 tok = AutoTokenizer.from_pretrained(settings.GEN_MODEL)
 model = AutoModelForCausalLM.from_pretrained(
     settings.GEN_MODEL,
@@ -30,15 +30,12 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto",
     low_cpu_mem_usage=True
 )
-
-gen = pipeline("text-generation", model=model, tokenizer=tok, max_new_tokens=800)
+gen = pipeline("text-generation", model=model, tokenizer=tok)
 
 def embed_texts(texts: list[str]):
-    v = emb_model.encode(texts, normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=False)
-    return v
+    return emb_model.encode(texts, normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=False)
 
 def add_to_index(records: list[dict]):
-    # records: [{"text":..., "source_url":..., "title":..., "doc_id":..., "chunk_id":...}]
     vecs = embed_texts([r["text"] for r in records])
     index.add(vecs)
     store.extend(records)
@@ -57,3 +54,13 @@ def search(query: str, top_k=8):
         rec["score"] = float(scores[0][rank])
         hits.append(rec)
     return hits
+
+def generate_chat(messages: list[dict], max_new_tokens=800, temperature=0.2):
+    """
+    messages: [{"role":"system"/"user"/"assistant", "content":"..."}]
+    """
+    prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    out = gen(prompt, do_sample=(temperature > 0.0), temperature=temperature, max_new_tokens=max_new_tokens)[0]["generated_text"]
+    # pipelineは入力+生成を返すことがあるため、プロンプトを取り除く
+    generated = out[len(prompt):].strip()
+    return generated or out