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

# Load your model and tokenizer using the adapter weights
model_name = "mherrador/CE5.0_expert_v2"  
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

# Explicitly set device to CPU
device = torch.device("cpu") 

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    # device_map="auto",  # Let Transformers choose the best device
    trust_remote_code=True,
).to(device)  # Move model to the specified device

tokenizer = AutoTokenizer.from_pretrained(model_name)

# Function to generate recommendations
def generate_recommendations(input_text):
    inputs = tokenizer(input_text, return_tensors="pt").to(device)  # Move input to device
    outputs = model.generate(**inputs, max_new_tokens=128)
    recommendations = tokenizer.batch_decode(outputs)[0]
    return recommendations

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_recommendations,
    inputs=gr.Textbox(lines=5, placeholder="Enter your questions here..."),
    outputs=gr.Textbox(lines=10),
    title="Circular Economy Recommender",
    description="Enter your questions about circular economy practices to get recommendations.",
)

# Launch the interface
iface.launch()