Airflow Tutorial (RAG 기반 로컬 AI 모델 활용)

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.utils.dates import days_ago
import datetime
import os
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

# DAG 설정
dag = DAG(
    dag_id='rag_local_ai_pipeline',
    default_args={
        'owner': 'airflow',
        'depends_on_past': False,
        'start_date': days_ago(2),
        'retries': 1,
        'retry_delay': datetime.timedelta(minutes=5),
    },
    schedule_interval=datetime.timedelta(days=1),
    catchup=False,
)

# 설정값
KNOWLEDGE_BASE_PATH = "/path/to/your/knowledge_base"  # 로컬 지식 베이스 경로
MODEL_NAME = "skt/kogpt2-base-v2"  # 사용할 KoGPT2 모델 이름
EMBEDDING_MODEL = "sentence-transformers/all-mpnet-base-v2" # 임베딩 모델 이름

# 1. 지식 베이스 업데이트 및 임베딩 생성 Task
def update_knowledge_base():
    """지식 베이스를 업데이트하고, 각 문서에 대한 임베딩을 생성하여 저장합니다."""
    print("Updating knowledge base...")
    documents = []
    for filename in os.listdir(KNOWLEDGE_BASE_PATH):
        if filename.endswith(".txt"):  # 예시: TXT 파일만 처리
            with open(os.path.join(KNOWLEDGE_BASE_PATH, filename), "r", encoding="utf-8") as f:
                documents.append(f.read())

    # 임베딩 모델 로드
    embedding_model = SentenceTransformer(EMBEDDING_MODEL)

    # 문서 임베딩 생성
    embeddings = embedding_model.encode(documents)

    # 임베딩 저장 (예시: numpy 배열로 저장)
    np.save("knowledge_base_embeddings.npy", embeddings)
    print("Knowledge base updated and embeddings generated.")
    return "knowledge_base_updated"

update_knowledge_base_task = PythonOperator(
    task_id='update_knowledge_base',
    python_callable=update_knowledge_base,
    dag=dag,
)


# 2. 사용자 질문 처리 및 답변 생성 Task
def generate_answer(ti=None):
    """사용자 질문을 받고, RAG 기반으로 답변을 생성합니다."""
    print("Generating answer...")
    data_updated = ti.xcom_pull(task_ids='update_knowledge_base', key='return_value')
    print(f"Knowledge base status: {data_updated}")

    # 사용자 질문 (예시)
    user_question = "Airflow는 무엇인가요?"

    # 임베딩 로드
    embeddings = np.load("knowledge_base_embeddings.npy")

    # 질문 임베딩 생성
    embedding_model = SentenceTransformer(EMBEDDING_MODEL)
    question_embedding = embedding_model.encode(user_question)

    # 코사인 유사도 계산
    similarities = cosine_similarity([question_embedding], embeddings)[0]

    # 가장 유사한 문서 선택
    most_similar_index = np.argmax(similarities)
    print(f"가장 유사한 문서 인덱스: {most_similar_index}")

    # 관련 문서 내용 가져오기 (실제로는 지식 베이스에서 가져와야 함)
    documents = []
    for filename in os.listdir(KNOWLEDGE_BASE_PATH):
        if filename.endswith(".txt"):
            with open(os.path.join(KNOWLEDGE_BASE_PATH, filename), "r", encoding="utf-8") as f:
                documents.append(f.read())
    relevant_document = documents[most_similar_index]

    # 모델 로드
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

    # RAG 기반 프롬프트 생성
    prompt = f"다음은 관련 문서입니다: {relevant_document}\n질문: {user_question}\n답변:"

    # 답변 생성
    generator = pipeline('text-generation', model=model, tokenizer=tokenizer, device=0) # GPU 사용 시 device=0, CPU 사용 시 device=-1
    generated_text = generator(prompt, max_length=200, num_return_sequences=1)[0]['generated_text']

    print(f"생성된 답변:\n{generated_text}")

generate_answer_task = PythonOperator(
    task_id='generate_answer',
    python_callable=generate_answer,
    dag=dag,
)

# Task 의존성 설정
update_knowledge_base_task >> generate_answer_task