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

MODEL_ID = "bharatgenai/AgriParam"

# Load tokenizer & model
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    trust_remote_code=True
).to("cpu")

def predict(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")

    # Generate multi-token text
    outputs = model.generate(
        **inputs,
        max_new_tokens=200,       # response length
        temperature=0.7,          # creativity
        top_p=0.9,                # nucleus sampling
        do_sample=True,            # sampling instead of greedy
        use_cache=False
    )

    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Remove prompt repetition
    return generated_text[len(prompt):].strip()

demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(lines=4, placeholder="Ask your agriculture question..."),
    outputs=gr.Textbox(label="AgriParam Answer"),
    title="AgriParam - Agriculture AI",
    description="Agriculture Q&A model fine-tuned for Indian farmers."
)

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