import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel

# Model setup
base_model = "HuggingFaceH4/zephyr-7b-beta"
model_path = r"hmorad/zephyr-geo-finetuned/tree/main/checkpoint-161"  # Update to your repo ID

# Quantization config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4"
)

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(base_model)
tokenizer.pad_token = tokenizer.eos_token

# Load base model
base = AutoModelForCausalLM.from_pretrained(
    base_model,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True
)

# Load fine-tuned LoRA adapter
model = PeftModel.from_pretrained(base, model_path)
model.eval()

def generate_response(prompt, max_new_tokens=200):
    input_text = f"user: {prompt}\nassistant:"
    inputs = tokenizer(input_text, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=0.6,
            do_sample=True,
            top_k=50,
            top_p=0.9,
            pad_token_id=tokenizer.eos_token_id
        )
    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
    return decoded.split("assistant:")[-1].strip()

# Gradio interface
iface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(lines=4, placeholder="Enter your geospatial query (e.g., 'Convert this shapefile to PostGIS')"),
    outputs="text",
    title="Geospatial Data Chatbot",
    description="Ask questions about geospatial data processing.",
    theme="huggingface"
)

if __name__ == "__main__":
    iface.launch()