281 lines
9.2 KiB
Python
281 lines
9.2 KiB
Python
from langchain_ibm import WatsonxEmbeddings
|
|
from langchain_ibm import ChatWatsonx
|
|
from langchain_core.prompts import ChatPromptTemplate
|
|
from langchain_core.output_parsers import StrOutputParser
|
|
from langchain_chroma import Chroma
|
|
from langchain_ollama import ChatOllama
|
|
from langchain_core.documents import Document
|
|
from langchain_core.messages import HumanMessage
|
|
|
|
import base64
|
|
import os
|
|
from pathlib import Path
|
|
from dotenv import load_dotenv
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
import shutil
|
|
|
|
#########
|
|
# 모델
|
|
#########
|
|
load_dotenv()
|
|
|
|
apikey = os.getenv("WATSONX_API_KEY")
|
|
project_id = os.getenv("WATSONX_PROJECT_ID")
|
|
watsonx_ai_url = os.getenv("WATSONX_URL")
|
|
|
|
watson_llm = ChatWatsonx(
|
|
model_id="ibm/granite-4-h-small",
|
|
url=f"{watsonx_ai_url}",
|
|
api_key=f"{apikey}",
|
|
project_id=f"{project_id}",
|
|
max_tokens=2000,
|
|
)
|
|
|
|
watson_embedding = WatsonxEmbeddings(
|
|
model_id="ibm/granite-embedding-278m-multilingual",
|
|
url=f"{watsonx_ai_url}",
|
|
api_key=f"{apikey}",
|
|
project_id=f"{project_id}"
|
|
)
|
|
|
|
vision_llm = ChatOllama(model="minimax-m3:cloud", temperature=0)
|
|
|
|
parser = StrOutputParser()
|
|
|
|
print("모델 초기화 완료")
|
|
print(f"LLM : ibm/granite-4-h-small")
|
|
print(f"Embedding : ibm/granite-embedding-278m-multilingual")
|
|
print(f"Vision LLM: minimax-m3:cloud")
|
|
|
|
|
|
#########
|
|
# 이미지 유틸리티 함수
|
|
#########
|
|
|
|
def encode_image(image_path:str)->str:
|
|
"""
|
|
이미지 파일을 Base64 문자열로 인코딩
|
|
"""
|
|
path = Path(image_path)
|
|
|
|
if not path.exists():
|
|
raise FileNotFoundError(f"파일이 존재하지 않습니다. : {image_path}")
|
|
|
|
with open(image_path, "rb") as f:
|
|
return base64.b64encode(f.read()).decode("utf-8")
|
|
|
|
#########
|
|
# 샘플 테스트 데이터
|
|
#########
|
|
|
|
def create_sample_images(image_dir = "./sample_images"):
|
|
os.makedirs(image_dir, exist_ok =True)
|
|
|
|
# 기본 이미지 생성
|
|
img1 = Image.new("RGB", (600, 400), color = (255, 255, 255))
|
|
draw = ImageDraw.Draw(img1)
|
|
|
|
draw.text((180, 20), "2024 분기별 매출 현황 (억원)", fill=(30, 30, 30))
|
|
|
|
bars = [
|
|
("Q1", 120, (70, 130, 180)),
|
|
("Q2", 185, (60, 179, 113)),
|
|
("Q3", 160, (255, 165, 0)),
|
|
("Q4", 230, (255, 50, 50)),
|
|
]
|
|
bar_w, base_y = 80, 330
|
|
|
|
for i, (label, val, color) in enumerate(bars):
|
|
x = 80 + i * 130
|
|
h = val
|
|
draw.rectangle([x, base_y - h, x + bar_w, base_y], fill=color)
|
|
draw.text((x + 20, base_y - h - 20), f"{val}억", fill=(30, 30, 30))
|
|
draw.text((x+25, base_y + 5), label, fill=(30, 30, 30))
|
|
draw.line([(60, 50), (60, base_y)], fill=(0, 0, 0), width=2)
|
|
draw.line([(60, base_y), (500, base_y)], fill=(0, 0, 0), width=2)
|
|
draw.text((10, 15), "상승 추세 : Q4 최고치 달성", fill=(0,0,0))
|
|
|
|
img1.save(f"{image_dir}/sales_chart_2024.jpg")
|
|
print(f"생성 : {image_dir}/sales_chart_2024.jpg")
|
|
|
|
return image_dir
|
|
|
|
def create_sample_text_docs():
|
|
"""테스트용 샘플 텍스트 문서 생성"""
|
|
docs = [
|
|
Document(
|
|
page_content="""
|
|
2024년 연간 매출 보고서 요약
|
|
|
|
당사의 2024년 전체 매출은 695억원으로, 전년 대비 23% 성장을 달성했습니다.
|
|
1분기(Q1) 120억원, 2분기(Q2) 185억원, 3분기(Q3) 160억원, 4분기(Q4) 230억원
|
|
기록했으며, Q4에 최고치를 달성했습니다.
|
|
|
|
주요 성장요인:
|
|
- 신제품 Pro-XL 출시로 인한 프리미엄 라인 매출 확대
|
|
- 해외 수출 비중 15% ~ 28% 증가
|
|
- 온라인 채널 전환으로 마진율 개선
|
|
|
|
4분기 실적은 연말 프로모션과 신규 파트너십 체결 효과로 큰 폭 상승했습니다.
|
|
""",
|
|
metadata = {
|
|
"source": "annual_report_2024.txt",
|
|
"type": "text",
|
|
"category": "재무"
|
|
}
|
|
)
|
|
]
|
|
|
|
return docs
|
|
|
|
def image_to_caption(image_path):
|
|
"""
|
|
이미지를 설명 텍스트로 변환
|
|
Vision LLM이 이미지를 분석하여 검색 간으한 텍스트 설명 생성
|
|
"""
|
|
img_b64 = encode_image(image_path)
|
|
message = HumanMessage(
|
|
content = [
|
|
{"type":"image_url", "image_url":{'url':f'data:image/jpeg;base64,{img_b64}'}},
|
|
{"type":"text", "text": """
|
|
이 이미지를 검색용 설명문으로 요약하세요.
|
|
200자 이내로 작성하세요.
|
|
핵심 객체와 텍스트만 포함하세요.
|
|
"""}
|
|
]
|
|
)
|
|
return vision_llm.invoke([message]).content
|
|
|
|
def build_multimodal_index(image_dir:str, text_docs:list[Document]):
|
|
"""
|
|
이미지 폴더의 모든 이미지들을 캡셔닝하여 텍스트 문서와 함께 벡터 인덱스 구축
|
|
"""
|
|
all_docs = list(text_docs)
|
|
|
|
# image_dir 안 파일 가져오기
|
|
# image_files = list(Path(image_dir).glob("*.{jpg, jpeg, png}"))
|
|
valid_extensions = {".jpg", ".jpeg", ".png", ".webp", ".JPG", ".JPEG", ".PNG"}
|
|
image_files = [
|
|
p for p in Path(image_dir).rglob("*")
|
|
if p.suffix in valid_extensions
|
|
]
|
|
|
|
if not image_files:
|
|
print(f"{image_dir}에 이미지를 찾지 못했습니다..")
|
|
else:
|
|
print(f"이미지 캡셔닝 {len(image_files)}개")
|
|
|
|
# 캡션 생성 => Document => 임베딩
|
|
for img_path in image_files:
|
|
caption = image_to_caption(image_path=img_path)
|
|
|
|
doc = Document(page_content=caption, metadata={
|
|
"source": str(img_path),
|
|
"type": "image",
|
|
"image_path": str(img_path)
|
|
})
|
|
all_docs.append(doc)
|
|
|
|
# 기존 DB 삭제
|
|
db_path = "./db/multimodal_db"
|
|
if Path(db_path).exists():
|
|
shutil.rmtree(db_path)
|
|
|
|
return Chroma.from_documents(all_docs, watson_embedding, persist_directory=db_path)
|
|
|
|
def search_with_images(vectorstore, query):
|
|
"""
|
|
벡터 검색 결과를 텍스트/이미지 결과로 분리하여 반환
|
|
"""
|
|
results = vectorstore.similarity_search_with_relevance_scores(query, k=5)
|
|
text_results = [doc for doc, score in results if doc.metadata.get("type") != 'image']
|
|
image_results = [doc for doc, score in results if doc.metadata.get("type") == 'image']
|
|
print(f"텍스트 결과 : {len(text_results)}, 이미지 결과 {len(image_results)}")
|
|
return text_results, image_results
|
|
|
|
def multimodal_answer(vectorstore, question):
|
|
"""
|
|
특스트 + 이미지 검색 결과를 통합하여 멀티모달 최종 답변 생성
|
|
args:
|
|
vectorstore: Chroma 벡터 스토어
|
|
question: 사용자 질문
|
|
|
|
return "
|
|
{"answer" : str, "images": list[srt]}
|
|
"""
|
|
text_results, image_results = search_with_images(vectorstore = vectorstore, query = question)
|
|
|
|
# 텍스트 결과 하나의 컨텍스트로 생성
|
|
text_context = "\n\n".join([r.page_content for r in text_results])
|
|
|
|
# 이미지로 검색된 경우
|
|
image_context = ""
|
|
referenced_images = []
|
|
for img_doc in image_results[:3]:
|
|
img_path = img_doc.metadata.get("image_path")
|
|
|
|
img_b64 = encode_image(img_path)
|
|
analysis = vision_llm.invoke([HumanMessage(
|
|
content = [
|
|
{"type":"image_url", "image_url":{'url':f'data:image/jpeg;base64,{img_b64}'}},
|
|
{"type":"text", "text": f"이 이미지에서 다음 질문과 관련된 내용을 설명하세요: {question}"}
|
|
])
|
|
]).content
|
|
|
|
image_context += f"[이미지 분석: {img_path}]\n{analysis}\n\n"
|
|
referenced_images.append(img_path)
|
|
|
|
# 최종 답변
|
|
combined_context = text_context + "\n\n" + image_context
|
|
|
|
final_prompt = ChatPromptTemplate.from_messages([
|
|
("system", "다음은 텍스트와 이미지 분석 결과를 참고하여 질문에 대한 답변입니다.\n\n{context}"),
|
|
("human", "{question}")
|
|
])
|
|
|
|
parser = StrOutputParser()
|
|
chain = final_prompt | watson_llm | parser
|
|
|
|
answer = chain.invoke({
|
|
"context": combined_context,
|
|
"question": question
|
|
})
|
|
|
|
return {"answer" : answer, "images": referenced_images}
|
|
|
|
|
|
#########
|
|
# 기존 벡터스토어 로드 함수
|
|
#########
|
|
|
|
def load_or_build_vectorstore(image_dir, text_docs, db_path="./db/multimodal_db", force_rebuild=False):
|
|
"""
|
|
이미 벡터스토어가 존재하면 로드하고, 없으면 새로 구축
|
|
"""
|
|
if not force_rebuild and Path(db_path).exists():
|
|
print(f"기존 멀티모달 DB 로드 완료: {db_path}")
|
|
return Chroma(embedding_function=watson_embedding, persist_directory=db_path)
|
|
|
|
return build_multimodal_index(image_dir, text_docs)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
IMAGE_DIR = "./sample_images"
|
|
# 샘플 이미지 생성
|
|
create_sample_images(IMAGE_DIR)
|
|
# 샘플 텍스트 생성
|
|
text_docs = create_sample_text_docs()
|
|
# 벡터스토어(인덱스) 생성
|
|
vectorstore = load_or_build_vectorstore(image_dir=IMAGE_DIR, text_docs=text_docs, force_rebuild=False)
|
|
# vectorstore = load_or_build_vectorstore(IMAGE_DIR, text_docs, "./db", False)
|
|
|
|
print("\n")
|
|
print("=" * 60)
|
|
print("질의 응답 테스트")
|
|
print("=" * 60)
|
|
result = multimodal_answer(vectorstore, "매출 차트의 추세는?")
|
|
print(result["answer"])
|
|
|
|
if result["images"]:
|
|
for img_path in result["images"]:
|
|
print(img_path) |