import gradio as gr
from transformers import AutoConfig, AutoModel
import torch

# Load the model
config = AutoConfig.from_pretrained("GodfreyOwino/NPK_prediction_model2", trust_remote_code=True)
model = AutoModel.from_pretrained("GodfreyOwino/NPK_prediction_model2", config=config, trust_remote_code=True)

def predict(crop_name, target_yield, field_size, ph, organic_carbon, nitrogen, phosphorus, potassium, soil_moisture):
    input_data = {
        'crop_name': [crop_name],
        'target_yield': [target_yield],
        'field_size': [field_size],
        'ph': [ph],
        'organic_carbon': [organic_carbon],
        'nitrogen': [nitrogen],
        'phosphorus': [phosphorus],
        'potassium': [potassium],
        'soil_moisture': [soil_moisture]
    }
    
    # Convert input data to tensors
    input_tensors = {k: torch.tensor(v) for k, v in input_data.items()}
    
    # Make prediction
    with torch.no_grad():
        prediction = model(input_tensors)
    
    # Convert prediction to a list if it's a tensor
    result = prediction.tolist() if isinstance(prediction, torch.Tensor) else prediction
    
    return str(result)  # Convert to string for Gradio output

# Define Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Textbox(label="Crop Name"),
        gr.Number(label="Target Yield"),
        gr.Number(label="Field Size"),
        gr.Number(label="pH"),
        gr.Number(label="Organic Carbon"),
        gr.Number(label="Nitrogen"),
        gr.Number(label="Phosphorus"),
        gr.Number(label="Potassium"),
        gr.Number(label="Soil Moisture")
    ],
    outputs="text",
    title="NPK Prediction Model",
    description="Enter the details to get NPK predictions."
)

iface.launch()