146 lines
5.2 KiB
Python
146 lines
5.2 KiB
Python
import os
|
|
import torch
|
|
from sentence_transformers import SentenceTransformer
|
|
import chromadb
|
|
from chromadb.utils import embedding_functions
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM
|
|
from fastapi import FastAPI
|
|
|
|
# # === 경로 설정 (모두 로컬) ===
|
|
QWEN_MODEL_PATH = "./models/Qwen3-0.6B"
|
|
EMBEDDING_MODEL_PATH = "./models/all-MiniLM-L6-v2"
|
|
|
|
# 2. 벡터 DB 설정
|
|
persist_directory = "./chroma_db"
|
|
chroma_client = chromadb.PersistentClient(path=persist_directory)
|
|
|
|
collection = chroma_client.get_or_create_collection(
|
|
name="manuals",
|
|
)
|
|
|
|
_model = None
|
|
_tokenizer = None
|
|
|
|
def get_qwen_model():
|
|
global _model, _tokenizer
|
|
if _model is not None:
|
|
return _model, _tokenizer
|
|
_tokenizer = AutoTokenizer.from_pretrained(
|
|
QWEN_MODEL_PATH,
|
|
trust_remote_code=True,
|
|
local_files_only=True
|
|
)
|
|
_model = AutoModelForCausalLM.from_pretrained(
|
|
QWEN_MODEL_PATH,
|
|
dtype=torch.bfloat16,
|
|
device_map="auto",
|
|
trust_remote_code=True,
|
|
local_files_only=True
|
|
)
|
|
# ✅ torch.compile() 적용 (PyTorch 2.0+)
|
|
if hasattr(torch, 'compile'):
|
|
try:
|
|
print("🚀 torch.compile() 적용 중...")
|
|
# mode="reduce-overhead": 추론 시 추천
|
|
# dynamic=False: 고정 입력 크기 시 더 빠름 (보통 RAG는 입력 길이가 유동적이므로 dynamic=True 권장)
|
|
_model = torch.compile(
|
|
_model,
|
|
mode="reduce-overhead", # 또는 "max-autotune" (초기 컴파일 시간 ↑, 이후 속도 ↑↑)
|
|
dynamic=True # 입력 길이가 매번 다르므로 동적 크기 허용
|
|
)
|
|
print("✅ torch.compile() 성공!")
|
|
except Exception as e:
|
|
print(f"⚠️ torch.compile() 실패, 원본 모델 사용: {e}")
|
|
pass # 실패하면 그냥 원본 사용
|
|
return _model, _tokenizer
|
|
|
|
def query_and_summarize(job: str, query: str, top_k: int = 3, min_similarity: float = 0.2):
|
|
from datetime import datetime
|
|
print(f'1 {datetime.now()}')
|
|
# 관련 문서 검색 (top_k보다 여유 있게 가져옴)
|
|
results = collection.query(
|
|
query_texts=[query],
|
|
n_results=top_k + 2, # 여유분 확보
|
|
where={"dept": job}
|
|
)
|
|
print(f'{datetime.now()}')
|
|
if not results['documents'][0]:
|
|
return "관련 문서를 찾을 수 없습니다."
|
|
|
|
print(f'2 {datetime.now()}')
|
|
# 유사도 계산 및 필터링
|
|
filtered_docs = []
|
|
for doc, dist in zip(results['documents'][0], results['distances'][0]):
|
|
similarity = 1 - dist
|
|
if similarity >= min_similarity:
|
|
filtered_docs.append((doc, similarity))
|
|
if len(filtered_docs) >= top_k:
|
|
break
|
|
print(f'3 {datetime.now()}')
|
|
if not filtered_docs:
|
|
return "유사도 기준을 만족하는 문서가 없습니다."
|
|
|
|
# 컨텍스트 생성 (유사도 내림차순 정렬은 이미 Chroma가 보장)
|
|
context_parts = []
|
|
for i, (doc, sim) in enumerate(filtered_docs):
|
|
context_parts.append(f"[청크 {i+1} | 유사도: {sim:.3f}]\n{doc}")
|
|
context = "\n\n".join(context_parts)
|
|
print(f'4 {datetime.now()}')
|
|
# 모델 로드
|
|
model, tokenizer = get_qwen_model()
|
|
print(f'5 {datetime.now()}')
|
|
# 프롬프트 구성
|
|
messages = [
|
|
{
|
|
"role": "system",
|
|
"content": (
|
|
"당신은 회사 재무/회계 업무 전문 어시스턴트입니다. "
|
|
"사용자에게 제공된 여러 청크를 종합하여, 정확하고 상세하게 답변하세요. "
|
|
"필요시 문서 내용을 직접 인용하거나 요약해도 됩니다. "
|
|
"추측하지 말고, 문서에 근거한 정보만 사용하세요."
|
|
)
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": f"다음 문서들을 참고하세요:\n\n{context}\n\n질문: {query}"
|
|
}
|
|
]
|
|
print(f'6 {datetime.now()}')
|
|
# 토큰화 및 생성
|
|
text = tokenizer.apply_chat_template(
|
|
messages,
|
|
tokenize=False,
|
|
add_generation_prompt=True,
|
|
enable_thinking=False # 작은 모델에선 비활성화 권장
|
|
)
|
|
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
|
|
print(f'7 {datetime.now()}')
|
|
generated_ids = model.generate(
|
|
**model_inputs,
|
|
max_new_tokens=150,
|
|
do_sample=True, # ✅ 샘플링 활성화
|
|
temperature=0.3,
|
|
top_p=0.9,
|
|
pad_token_id=tokenizer.eos_token_id
|
|
)
|
|
print(f'8 {datetime.now()}')
|
|
input_len = model_inputs.input_ids.shape[1]
|
|
output_ids = generated_ids[0][input_len:]
|
|
response = tokenizer.decode(output_ids, skip_special_tokens=True).strip()
|
|
print(f'9 {datetime.now()}')
|
|
return response
|
|
|
|
# FastAPI 앱
|
|
app = FastAPI()
|
|
|
|
@app.get("/")
|
|
def question(query: str):
|
|
answer = query_and_summarize(job="FI", query=query)
|
|
return {"answer": answer}
|
|
|
|
# 개발용 실행 (직접 실행 시)
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
print("서버 시작: uvicorn manual:app --reload")
|
|
# 예시 질의 (주석 해제 시 직접 테스트 가능)
|
|
# print(query_and_summarize("FI", "외화 송금 절차는 어떻게 되나요?")) |