pip install smolagents pandas langchain langchain-community sentence-transformers datasets python-dotenv rank_bm25 --upgrade

# 환경 변수 로드 (HF_TOKEN 포함)
from dotenv import load_dotenv
load_dotenv()

import datasets
from langchain.docstore.document import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.retrievers import BM25Retriever

# 허깅 페이스 설명서 데이터 세트 로드
knowledge_base = datasets.load_dataset("m-ric/huggingface_doc", split="train")

# Transformers 라이브러리 설명서만 포함하도록 필터링
knowledge_base = knowledge_base.filter(lambda row: row["source"].startswith("huggingface/transformers"))

# 데이터 세트 항목을 메타데이터가 있는 Document 객체로 변환
source_docs = [
    Document(page_content=doc["text"], metadata={"source": doc["source"].split("/")[1]})
    for doc in knowledge_base
]

# 더 나은 검색을 위해 문서를 작은 청크로 분할
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,  # 청크당 문자 수
    chunk_overlap=50,  # 컨텍스트 유지를 위한 청크 간 중첩
    add_start_index=True,
    strip_whitespace=True,
    separators=["\n\n", "\n", ".", " ", ""],  # 분할 우선순위
)
docs_processed = text_splitter.split_documents(source_docs)

print(f"Knowledge base prepared with {len(docs_processed)} document chunks")

from smolagents import Tool

class RetrieverTool(Tool):
    name = "retriever"
    description = "의미 기반 검색을 사용하여 쿼리에 답변하는 데 가장 관련성이 높은 transformers 설명서 부분을 검색합니다."
    inputs = {
        "query": {
            "type": "string",
            "description": "수행할 쿼리입니다. 대상 문서와 의미적으로 가까워야 합니다. 질문보다는 긍정문을 사용하세요.",
        }
    }
    output_type = "string"

    def __init__(self, docs, **kwargs):
        super().__init__(**kwargs)
        # 처리된 문서로 검색기 초기화
        self.retriever = BM25Retriever.from_documents(
            docs, k=10  # 가장 관련성 높은 상위 10개 문서 반환
        )

    def forward(self, query: str) -> str:
        """제공된 쿼리를 기반으로 검색을 실행합니다."""
        assert isinstance(query, str), "Your search query must be a string"

        # 관련 문서 검색
        docs = self.retriever.invoke(query)

        # 가독성을 위해 검색된 문서 형식 지정
        return "\nRetrieved documents:\n" + "".join(
            [
                f"\n\n===== Document {str(i)} =====\n" + doc.page_content
                for i, doc in enumerate(docs)
            ]
        )

# 처리된 문서로 검색 도구 초기화
retriever_tool = RetrieverTool(docs_processed)

from smolagents import InferenceClientModel, CodeAgent

# 검색 도구로 에이전트 초기화
agent = CodeAgent(
    tools=[retriever_tool],  # 에이전트가 사용할 수 있는 도구 목록
    model=InferenceClientModel(),  # 기본 모델 "Qwen/Qwen2.5-Coder-32B-Instruct"
    max_steps=4,  # 논리적 추론 단계 수 제한
    verbosity_level=2,  # 에이전트의 상세한 논리적 추론 과정 표시
)

# 특정 모델을 사용하려면 다음과 같이 지정할 수 있습니다:
# model=InferenceClientModel(model_id="meta-llama/Llama-3.3-70B-Instruct")

# 정보 검색이 필요한 질문하기
question = "For a transformers model training, which is slower, the forward or the backward pass?"

# 에이전트를 실행하여 답변 얻기
agent_output = agent.run(question)

# 최종 답변 표시
print("\nFinal answer:")
print(agent_output)