app.py

import gradio as gr
from llama_index.core.llms import ChatMessage
from llama_index.core.chat_engine.condense_plus_context import CondensePlusContextChatEngine
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
from llama_index.core.node_parser import SentenceSplitter
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core import Settings
from llama_cpp import Llama

# Load Llama model
llm = Llama(
    model_path="zephyr-7b-beta.Q4_K_M.gguf",  # Sesuaikan path model
    n_gpu_layers=1,  # Aktifkan GPU jika tersedia
)

# Konfigurasi untuk indexing dokumen
documents = SimpleDirectoryReader(input_files=[    "bahandokumen/bonus.txt",
                                                   "bahandokumen/cuti.txt",
                                                   "bahandokumen/disiplinkerja.txt",
                                                   "bahandokumen/fasilitas&bantuan.txt",
                                                   "bahandokumen/fasilitaskerja.txt",
                                                   "bahandokumen/hak.txt",
                                                   "bahandokumen/hubunganpengusaha&serikat.txt",
                                                   "bahandokumen/istilah.txt",
                                                   "bahandokumen/jaminanserikat.txt",
                                                   "bahandokumen/jamkes.txt",
                                                   "bahandokumen/jamsos.txt",
                                                   "bahandokumen/keluhkesah.txt",
                                                   "bahandokumen/kenaikanupah.txt",
                                                   "bahandokumen/kewajiban.txt",
                                                   "bahandokumen/kompensasi.txt",
                                                   "bahandokumen/larangan.txt",
                                                   "bahandokumen/lembur.txt",
                                                   "bahandokumen/luaskesepakatan.txt",
                                                   "bahandokumen/mogok.txt",
                                                   "bahandokumen/pelanggaran&sanksi.txt",
                                                   "bahandokumen/pendidikan.txt",
                                                   "bahandokumen/pengangkatan.txt",
                                                   "bahandokumen/penilaian&promosi.txt",
                                                   "bahandokumen/pensiun.txt",
                                                   "bahandokumen/perjadin.txt",
                                                   "bahandokumen/pesangon.txt",
                                                   "bahandokumen/phk.txt",
                                                   "bahandokumen/pihak.txt",
                                                   "bahandokumen/pkb.txt",
                                                   "bahandokumen/resign.txt",
                                                   "bahandokumen/sanksi.txt",
                                                   "bahandokumen/shift.txt",
                                                   "bahandokumen/syaratkerja.txt",
                                                   "bahandokumen/tatacara.txt",
                                                   "bahandokumen/tka.txt",
                                                   "bahandokumen/tunjangan.txt",
                                                   "bahandokumen/uangpisah.txt",
                                                   "bahandokumen/upah.txt",
                                                   "bahandokumen/upahlembur.txt",
                                                   "bahandokumen/waktukerja.txt"]).load_data()

# Split dokumen menjadi node
parser = SentenceSplitter(chunk_size=300, chunk_overlap=20)
nodes = parser.get_nodes_from_documents(documents)

# Set embeddings dan index
embedding = HuggingFaceEmbedding("BAAI/bge-base-en-v1.5")
Settings.llm = llm
Settings.embed_model = embedding

index = VectorStoreIndex(nodes)

# Retrieve data berdasarkan query
retriever = index.as_retriever(similarity_top_k=3)
query_engine = index.as_query_engine(similarity_top_k=3)

# Setup chat engine
chat_engine = CondensePlusContextChatEngine.from_defaults(
    retriever=retriever,
    verbose=True,
)

# Fungsi untuk merespons chat
def generate_response(message, history):
    chat_messages = [
        ChatMessage(
            role="system",
            content="Anda adalah chatbot yang selalu menjawab pertanyaan dalam bahasa Indonesia dengan singkat dan ramah."
        )
    ]

    for human, ai in history:
        chat_messages.append(ChatMessage(role="user", content=human))
        chat_messages.append(ChatMessage(role="assistant", content=ai))

    chat_messages.append(ChatMessage(role="user", content=message))

    response = chat_engine.stream_chat(chat_messages)
    text = ""
    for token in response.response_gen:
        text += token
        yield text

# Fungsi untuk mereset riwayat chat
def clear_history():
    chat_engine.reset()

# Gradio UI
with gr.Blocks() as demo:
    with gr.Row():
        clear_btn = gr.Button("Clear")
        clear_btn.click(clear_history)

    chat_interface = gr.ChatInterface(
        fn=generate_response, 
        clear_btn=clear_btn,
        chatbot_style="default"
    )

demo.launch()