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

# Load models and tokenizers
base_model_name = "unsloth/Llama-3.2-1B-Instruct"
finetuned_model_name = "exo-is/esg-context-llama-1Bst-11M"  # Replace with your model's path

tokenizer = AutoTokenizer.from_pretrained(base_model_name)
base_model = AutoModelForCausalLM.from_pretrained(base_model_name, torch_dtype=torch.float32, low_cpu_mem_usage=True, device_map="cpu")
finetuned_model = AutoModelForCausalLM.from_pretrained(finetuned_model_name, torch_dtype=torch.float32, low_cpu_mem_usage=True, device_map="cpu")

def generate_text(model, prompt, max_new_tokens, temperature):
    inputs = tokenizer(prompt, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            temperature=temperature,
            num_return_sequences=1,
            do_sample=True,
        )
    
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

def log_interaction(model, prompt, output, validation):
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    with open('interaction_log.csv', 'a', newline='') as file:
        writer = csv.writer(file)
        writer.writerow([timestamp, model, prompt, output, validation])

def generate_and_compare(prompt, max_new_tokens, temperature):
    base_output = generate_text(base_model, prompt, max_new_tokens, temperature)
    finetuned_output = generate_text(finetuned_model, prompt, max_new_tokens, temperature)
    return base_output, finetuned_output

def evaluate_base(output, score):
    log_interaction("Base", gr.get_state('last_prompt'), output, score)
    return f"Avaliação registrada para o modelo Base: {score}"

def evaluate_finetuned(output, score):
    log_interaction("Fine-tuned", gr.get_state('last_prompt'), output, score)
    return f"Avaliação registrada para o modelo Fine-tuned: {score}"

with gr.Blocks() as demo:
    gr.Markdown("# Comparação de Modelos: Llama-3.2-1B-Instruct vs. Modelo Fine-tuned para Sustentabilidade")
    
    with gr.Row():
        prompt = gr.Textbox(lines=5, label="Insira seu prompt aqui")
        max_new_tokens = gr.Slider(50, 500, value=200, step=1, label="Máximo de Novos Tokens")
        temperature = gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperatura")
    
    generate_btn = gr.Button("Gerar")
    
    with gr.Row():
        with gr.Column():
            base_output = gr.Textbox(label="Saída do Modelo Base", lines=10)
            base_rating = gr.Radio(["1", "2", "3", "4", "5"], label="Avalie a resposta do Modelo Base")
            base_submit = gr.Button("Enviar Avaliação (Base)")
        
        with gr.Column():
            finetuned_output = gr.Textbox(label="Saída do Modelo Fine-tuned", lines=10)
            finetuned_rating = gr.Radio(["1", "2", "3", "4", "5"], label="Avalie a resposta do Modelo Fine-tuned")
            finetuned_submit = gr.Button("Enviar Avaliação (Fine-tuned)")
    
    base_feedback = gr.Textbox(label="Feedback da Avaliação (Base)")
    finetuned_feedback = gr.Textbox(label="Feedback da Avaliação (Fine-tuned)")
    
    generate_btn.click(generate_and_compare, inputs=[prompt, max_new_tokens, temperature], outputs=[base_output, finetuned_output])
    base_submit.click(evaluate_base, inputs=[base_output, base_rating], outputs=base_feedback)
    finetuned_submit.click(evaluate_finetuned, inputs=[finetuned_output, finetuned_rating], outputs=finetuned_feedback)
    
    demo.load(lambda: gr.update(value=""), outputs=[prompt])
    prompt.change(lambda x: gr.set_state(last_prompt=x), inputs=[prompt])

demo.launch()