app.py

import gradio as gr
import gspread
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from oauth2client.service_account import ServiceAccountCredentials

# =============== 1. Cache Google Sheets ===============
cached_text_data = None

def read_google_sheets():
    global cached_text_data
    if cached_text_data is not None:
        return cached_text_data

    try:
        scope = ["https://www.googleapis.com/auth/spreadsheets", "https://www.googleapis.com/auth/drive"]
        creds = ServiceAccountCredentials.from_json_keyfile_name("credentials.json", scope)
        client = gspread.authorize(creds)

        SPREADSHEET_ID = "1e_cNMhwF-QYpyYUpqQh-XCw-OdhWS6EuYsoBUsVtdNg"
        sheet_names = ["datatarget", "datacuti", "dataabsen", "datalembur", "pkb"]

        all_data = []
        spreadsheet = client.open_by_key(SPREADSHEET_ID)

        for sheet_name in sheet_names:
            try:
                sheet = spreadsheet.worksheet(sheet_name)
                data = sheet.get_all_values()
                all_data.append(f"=== Data dari {sheet_name.upper()} ===")
                all_data.extend([" | ".join(row) for row in data])
                all_data.append("\n")
            except gspread.exceptions.WorksheetNotFound:
                all_data.append(f"❌ Worksheet {sheet_name} tidak ditemukan.")

        cached_text_data = "\n".join(all_data).strip()
        return cached_text_data

    except Exception as e:
        return f"❌ ERROR: {str(e)}"


# =============== 2. Load Model Transformers ===============
def load_model():
    model_id = "HuggingFaceH4/zephyr-7b-beta"  # Bisa ganti ke Zephyr juga
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        device_map="auto",  # 🧠 Otomatis ke GPU jika tersedia
        torch_dtype=torch.float16
    )
    pipe = pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer,
        device=0 if torch.cuda.is_available() else -1,
        max_new_tokens=512,
        temperature=0.7,
        repetition_penalty=1.2,
        do_sample=True,
    )
    return pipe


# =============== 3. Buat Prompt dan Jawaban ===============
def generate_prompt(user_message, context_data):
    prompt = f"""
### SISTEM:
Anda adalah asisten HRD yang membantu karyawan memahami administrasi perusahaan. Jawablah dengan Bahasa Indonesia yang profesional dan ramah. Jika tidak tahu, katakan tidak tahu dengan sopan. Jangan jawab dalam bahasa Inggris.

### DATA:
{context_data}

### PERTANYAAN:
{user_message}

### JAWABAN:
"""
    return prompt.strip()


# =============== 4. Generate Response ===============
def generate_response(message, history, pipe):
    context = read_google_sheets()
    full_prompt = generate_prompt(message, context)
    response = pipe(full_prompt)[0]["generated_text"]

    # Bersihkan jawaban: ambil bagian setelah ### JAWABAN:
    cleaned = response.split("### JAWABAN:")[-1].strip()
    history = history or []
    history.append((message, cleaned))
    return cleaned


# =============== 5. Jalankan Gradio ===============
def main():
    pipe = load_model()

    def chatbot_response(message, history):
        return generate_response(message, history, pipe)

    gr.Interface(
        fn=chatbot_response,
        inputs=["text"],
        outputs=["text"],
        title="Chatbot HRD - Transformers",
        theme="compact"
    ).launch(share=True)

if __name__ == "__main__":
    main()