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

model_name = "bharatgenai/AgriParam"

print("Loading AgriParam model... This may take some time on free CPU")

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=False)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    trust_remote_code=True,
    torch_dtype=torch.float16,          # Changed to float16 (more compatible on CPU)
    device_map="cpu",
    low_cpu_mem_usage=True,
    attn_implementation="eager"         # This helps avoid rope_scaling issues
)

def chat(message, history):
    prompt = f"<user> {message} <assistant>"
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    outputs = model.generate(
        **inputs,
        max_new_tokens=400,
        temperature=0.7,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        use_cache=False                    # Important for stability on CPU
    )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    cleaned = response.split("<assistant>")[-1].strip()
    return cleaned

demo = gr.ChatInterface(
    fn=chat,
    title="🌾 AgriParam - Agriculture Assistant",
    description="Test version on free CPU • English, Hindi & Marathi",
    examples=[
        ["महाराष्ट्रात कापूस पिकावर लाल किडीचा उपाय काय आहे?"],
        ["Best practices for organic wheat farming in Maharashtra?"],
        ["उत्तर प्रदेश में गेहूं की खेती के लिए जैविक खाद कैसे बनाएं?"],
        ["महाराष्ट्रातील सोयाबीन पिकाला पिवळी पाने येण्याचे कारण आणि उपाय?"]
    ]
)

demo.launch()