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

# Base and adapter model paths
base_model = "HuggingFaceH4/zephyr-7b-beta"
adapter_path = "hmorad/zephyr-geo-finetuned"  # or use "hmorad/zephyr-geo-finetuned" if uploaded and public

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

# Load base model
base = AutoModelForCausalLM.from_pretrained(
    base_model,
    torch_dtype=torch.float32,  # Use float32 on CPU
    device_map="auto",          # This will place it correctly even on CPU
    trust_remote_code=True
)

# Load LoRA adapter
model = PeftModel.from_pretrained(base, adapter_path)
model.eval()

# Generation function
def generate_response(prompt, max_new_tokens=200):
    input_text = f"<|system|>\nYou are a helpful assistant for geospatial data.\n<|user|>\n{prompt}\n<|assistant|>\n"
    inputs = tokenizer(input_text, return_tensors="pt", padding=True).to("cpu")
    
    with torch.no_grad():
        output = model.generate(
            input_ids=inputs["input_ids"],
            attention_mask=inputs["attention_mask"],
            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)
    response = decoded.split("<|assistant|>")[-1].strip()
    return response

# 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')",
        label="Your Query"
    ),
    outputs=gr.Textbox(label="Response"),
    title="Geospatial Data Chatbot",
    description="Ask technical questions about raster/vector data, GIS tools, or formats like .tif/.shp/PostGIS.",
)

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