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

# Load your model and tokenizer using the adapter weights
model_name = "mherrador/CE-5.0"  # Replace with your actual model name
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Function to generate recommendations
def generate_recommendations(input_text):
    inputs = tokenizer(input_text, return_tensors="pt").to("cuda")
    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()