import gradio as gr
from ctransformers import AutoModelForCausalLM
import os

# Download model (runs on first launch)
MODEL_PATH = "Phi-3.5-mini-3.8B-ArliAI-RPMax-v1.1-Q6_K_L.gguf"
if not os.path.exists(MODEL_PATH):
    os.system(f"wget https://huggingface.co/ArliAI/Phi-3.5-mini-3.8B-ArliAI-RPMax-v1.1-GGUF/resolve/main/{MODEL_PATH}")

# Load GGUF model with ctransformers
llm = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    model_type="phi3",
    gpu_layers=50,  # Offload to GPU (set to 0 for CPU)
    context_length=2048
)

# System prompt to restrict knowledge
SYSTEM_PROMPT = """[SYSTEM] You are a compliance assistant. Follow these rules:
1. ONLY use data from '/data/company_policies.pdf' (provided in this Space's files)
2. If asked about unverified information, respond: "I can only reference approved documents"
3. Keep answers under 2 sentences."""

def respond(message, history):
    # Format Phi-3 prompt template
    prompt = f"{SYSTEM_PROMPT}\n[USER]{message}\n[ASSISTANT]"
    
    # Generate response
    response = llm(
        prompt,
        max_new_tokens=100,
        temperature=0.3,  # Low for deterministic answers
        stop=["[USER]", "\n\n"]
    )
    
    return response

# Gradio interface with file upload for knowledge base
with gr.Blocks() as demo:
    gr.Markdown("## Phi-3.5 Mini - Restricted Knowledge Assistant")
    with gr.Tab("Chat"):
        chat_interface = gr.ChatInterface(respond)
    with gr.Tab("Upload Source"):
        gr.File(label="Upload PDF/JSON for reference", file_count="single")
    
demo.launch()