Update app.py

app.py

@@ -1,13 +1,9 @@
-# main.py (Versi Final Definitif - Dengan Parsing Cerdas & Prompt Tepat)
+# app.py (Versi Final untuk Gradio di Hugging Face)
+
+import gradio as gr
 import os
-import shutil
 import re
-import uvicorn
-from fastapi import FastAPI, UploadFile, File, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-import torch
-
+import shutil
 from langchain_community.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.embeddings import HuggingFaceEmbeddings
@@ -15,128 +11,141 @@ from langchain_community.vectorstores import FAISS
 from langchain_community.retrievers import BM25Retriever
 from langchain.retrievers import EnsembleRetriever
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch._dynamo as dynamo
-
-dynamo.config.automatic_dynamic_shapes = False
-dynamo.config.assume_static_by_default = True
-
-UPLOAD_DIR = "temp_uploads"
-os.makedirs(UPLOAD_DIR, exist_ok=True)
-
-app = FastAPI(title="Financial RAG Chatbot API")
-
-origins = ["*"]
-app.add_middleware(CORSMiddleware, allow_origins=origins, allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
-
-rag_pipeline = {"retriever": None, "llm": None, "tokenizer": None, "embeddings": None, "all_chunks": None}
+import torch
 
-def setup_rag_pipeline():
-    print("Memulai setup RAG pipeline...")
+# --- 1. SETUP MODEL (dijalankan sekali saat aplikasi start) ---
+@torch.no_grad()
+def load_models():
+    print("Memuat model (hanya terjadi sekali)...")
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    print(f"Menggunakan device: {device}")
-    print("Memuat Embedding Model...")
-    rag_pipeline["embeddings"] = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", model_kwargs={'device': device})
-    print("Memuat LLM (Gemma 270M)...")
-    rag_pipeline["tokenizer"] = AutoTokenizer.from_pretrained("google/gemma-3-270m-it")
-    rag_pipeline["llm"] = AutoModelForCausalLM.from_pretrained("google/gemma-3-270m-it", device_map="auto", torch_dtype=torch.bfloat16)
-    print("Setup RAG pipeline selesai.")
-
-@app.on_event("startup")
-async def startup_event():
-    setup_rag_pipeline()
-
-@app.post("/upload")
-async def upload_document(file: UploadFile = File(...)):
-    # Fungsi ini tetap sama
+    
+    cache_dir = "./model_cache"
+    os.makedirs(cache_dir, exist_ok=True)
+    os.environ['SENTENCE_TRANSFORMERS_HOME'] = cache_dir
+
+    embeddings = HuggingFaceEmbeddings(
+        model_name="sentence-transformers/all-MiniLM-L6-v2",
+        cache_folder=cache_dir
+    )
+    
+    # Gunakan token dari secrets jika ada
+    hf_token = os.getenv("HUGGING_FACE_HUB_TOKEN")
+
+    tokenizer = AutoTokenizer.from_pretrained("google/gemma-3-270m-it", cache_dir=cache_dir, token=hf_token)
+    llm = AutoModelForCausalLM.from_pretrained(
+        "google/gemma-3-270m-it",
+        cache_dir=cache_dir,
+        device_map="auto",
+        torch_dtype=torch.bfloat16,
+        token=hf_token
+    )
+    print("Model berhasil dimuat.")
+    return embeddings, tokenizer, llm
+
+embeddings, tokenizer, llm = load_models()
+# Inisialisasi state global untuk retriever dan chunks
+rag_pipeline = {"retriever": None, "all_chunks": None}
+
+
+# --- 2. FUNGSI INTI RAG (backend logic) ---
+def process_document(uploaded_file):
+    if uploaded_file is None:
+        return "Mohon unggah file terlebih dahulu.", gr.update(interactive=False)
+        
     try:
-        file_path = os.path.join(UPLOAD_DIR, file.filename)
-        with open(file_path, "wb") as buffer: shutil.copyfileobj(file.file, buffer)
+        # Gradio menyimpan file di temporary path, kita bisa langsung pakai
+        file_path = uploaded_file.name
+        
         loader = PyPDFLoader(file_path)
         docs = loader.load()
         text_splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=200)
         chunks = text_splitter.split_documents(docs)
+        
         rag_pipeline["all_chunks"] = chunks
-        faiss_db = FAISS.from_documents(chunks, rag_pipeline["embeddings"])
+        
+        faiss_db = FAISS.from_documents(chunks, embeddings)
         faiss_retriever = faiss_db.as_retriever(search_kwargs={"k": 10})
         bm25_retriever = BM25Retriever.from_documents(chunks)
         bm25_retriever.k = 10
-        rag_pipeline["retriever"] = EnsembleRetriever(retrievers=[bm25_retriever, faiss_retriever], weights=[0.5, 0.5])
-        return {"status": "sukses", "message": f"Dokumen '{file.filename}' berhasil diproses."}
-    except Exception as e: raise HTTPException(status_code=500, detail=str(e))
-    finally:
-        if 'file_path' in locals() and os.path.exists(file_path): os.remove(file_path)
+        
+        rag_pipeline["retriever"] = EnsembleRetriever(
+            retrievers=[bm25_retriever, faiss_retriever],
+            weights=[0.5, 0.5]
+        )
+        
+        return f"File '{os.path.basename(file_path)}' berhasil diproses! Silakan ajukan pertanyaan.", gr.update(interactive=True)
+    except Exception as e:
+        return f"Error saat memproses file: {str(e)}", gr.update(interactive=False)
 
-class ChatRequest(BaseModel):
-    query: str
 
-@app.post("/chat")
-async def chat_with_doc(request: ChatRequest):
-    if not rag_pipeline.get("retriever"):
-        raise HTTPException(status_code=400, detail="Dokumen belum di-upload.")
+def get_rag_response(query, chat_history):
+    if rag_pipeline["retriever"] is None:
+        return "Dokumen belum diproses. Mohon unggah file terlebih dahulu."
 
-    query_original = request.query
+    query_original = query
     query_lower = query_original.lower()
-    clean_context = ""
-    found_source = ""
-    
+    final_answer = ""
+    found_source = "Tidak ada sumber spesifik"
+
     priority_keywords = ["jumlah aset lancar"]
     use_smart_lane = any(keyword in query_lower for keyword in priority_keywords)
 
     if use_smart_lane:
-        print("Jalur Cerdas Aktif! Mencari konteks & parsing tahun...")
+        # Jalur Cerdas
         year_match = re.search(r'\b(202[3-4])\b', query_lower)
         target_year = year_match.group(1) if year_match else "2024"
-        
-        all_chunks = rag_pipeline.get("all_chunks", [])
-        for chunk in all_chunks:
+        for chunk in rag_pipeline["all_chunks"]:
             lines = chunk.page_content.split('\n')
             for line in lines:
                 if any(keyword in line.lower() for keyword in priority_keywords):
-                    # --- [FIX #1] PARSING CERDAS DENGAN REGEX ---
-                    # Mencari semua angka yang diformat dengan koma/titik
                     numbers = re.findall(r'(\d{1,3}(?:[.,]\d{3})*)', line)
                     if len(numbers) >= 2:
-                        try:
-                            # Asumsi angka pertama adalah 2024, kedua 2023
-                            value_2024 = numbers[0]
-                            value_2023 = numbers[1]
-                            value = value_2024 if target_year == "2024" else value_2023
-                            # Buat fakta yang bersih untuk LLM
-                            clean_context = f"Fakta: jumlah aset lancar untuk tahun {target_year} adalah {value}."
-                            found_source = f"Halaman {chunk.metadata.get('page', 'NA')}"
-                            break
-                        except (IndexError, ValueError):
-                            continue
-            if clean_context:
-                break
+                        value_2024 = numbers[0]
+                        value_2023 = numbers[1]
+                        value = value_2024 if target_year == "2024" else value_2023
+                        final_answer = f"Jumlah aset lancar untuk tahun {target_year} adalah **{value}**."
+                        found_source = f"Sumber: Halaman {chunk.metadata.get('page', 'NA')}"
+                        break
+            if final_answer: break
     
-    if not clean_context:
-        print("Menggunakan Jalur Normal (Hybrid Search)...")
+    if not final_answer:
+        # Jalur Normal
         retrieved_docs = rag_pipeline["retriever"].invoke(query_original)
         clean_context = "\n\n".join([doc.page_content for doc in retrieved_docs[:3]])
         found_source = ", ".join(list(set([f"Halaman {doc.metadata.get('page', 'NA')}" for doc in retrieved_docs[:3]])))
-
-    # --- Generation ---
-    tokenizer = rag_pipeline["tokenizer"]
-    model = rag_pipeline["llm"]
-    
-    # --- [FIX #2] PROMPT YANG SANGAT TO-THE-POINT ---
-    chat_template = [
-        {"role": "user", "content": f"Gunakan informasi ini: '{clean_context}'. Jawab pertanyaan ini: '{query_original}'"}
-    ]
-    final_prompt = tokenizer.apply_chat_template(chat_template, tokenize=False, add_generation_prompt=True)
-    inputs = tokenizer(final_prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=150, do_sample=False, pad_token_id=tokenizer.eos_token_id)
+        
+        chat_template = [{"role": "system", "content": "Anda adalah AI analis keuangan yang teliti. Jawab pertanyaan hanya berdasarkan teks yang diberikan."}, {"role": "user", "content": f"Dari TEKS di bawah, temukan jawaban untuk pertanyaan '{query_original}'.\n\nTEKS:\n{clean_context}\n\nJAWABAN:"}]
+        final_prompt = tokenizer.apply_chat_template(chat_template, tokenize=False, add_generation_prompt=True)
+        inputs = tokenizer(final_prompt, return_tensors="pt").to(llm.device)
+        outputs = llm.generate(**inputs, max_new_tokens=250, do_sample=False, pad_token_id=tokenizer.eos_token_id)
+        input_length = inputs.input_ids.shape[1]
+        generated_tokens = outputs[0, input_length:]
+        final_answer = tokenizer.decode(generated_tokens, skip_special_tokens=True)
+
+    full_response = f"{final_answer}\n\n*{found_source}*"
+    chat_history.append((query, full_response))
+    return "", chat_history
+
+
+# --- 3. MEMBUAT UI DENGAN GRADIO ---
+with gr.Blocks() as demo:
+    gr.Markdown("# 📊 Financial RAG Chatbot")
     
-    input_length = inputs.input_ids.shape[1]
-    generated_tokens = outputs[0, input_length:]
-    final_answer = tokenizer.decode(generated_tokens, skip_special_tokens=True)
-
-    return {"answer": final_answer.strip(), "source": found_source}
-
-@app.get("/")
-def read_root():
-    return {"message": "Selamat datang di Financial RAG Chatbot API!"}
-
+    with gr.Row():
+        with gr.Column(scale=1):
+            file_output = gr.Textbox(label="Status Dokumen", interactive=False)
+            upload_button = gr.UploadButton("Klik untuk Upload PDF", file_types=[".pdf"])
+            ask_button = gr.Button("Tanya", interactive=False)
+
+        with gr.Column(scale=4):
+            chatbot = gr.Chatbot(label="Chat")
+            msg = gr.Textbox(label="Ketik Pertanyaan Anda di Sini...")
+            
+    # Hubungkan Aksi dengan Fungsi
+    upload_button.upload(process_document, upload_button, [file_output, ask_button])
+    msg.submit(get_rag_response, [msg, chatbot], [msg, chatbot])
+    ask_button.click(get_rag_response, [msg, chatbot], [msg, chatbot])
+
+# --- 4. JALANKAN APLIKASI ---
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    demo.launch()