import streamlit as st
from transformers import AutoTokenizer
from peft import AutoPeftModelForCausalLM
import torch

# Configuración del modelo y tokenizer
model_id = "TheBloke/Mistral-7B-Instruct-v0.2-GPTQ"
adapter = "nmarafo/Mistral-7B-Instruct-v0.2-TrueFalse-Feedback-GPTQ"
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True, return_token_type_ids=False)
tokenizer.pad_token = tokenizer.eos_token
model = AutoPeftModelForCausalLM.from_pretrained(model_id, adapter_name=adapter).cuda()

def generate_response(question, best_answer, student_answer):
    system_message = "Analyze the question, the expected answer, and the student's response. Determine if the student's answer is conceptually correct in relation to the expected answer, regardless of the exact wording. Return True if the student's answer is correct or False otherwise. Add a brief comment explaining the rationale behind the answer being correct or incorrect."
    prompt = f"{system_message}\n\nQuestion: {question}\nExpected Answer: {best_answer}\nStudent Answer: {student_answer}"
    prompt_template = f'<s>[INST] {prompt} [/INST]'
    input_ids = tokenizer(prompt_template, return_tensors='pt').input_ids.cuda()
    output = model.generate(input_ids, temperature=0.7, do_sample=True, top_p=0.95, top_k=40, max_new_tokens=512)
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    return response

# Crear la interfaz de usuario en Streamlit
st.title("Evaluador de Respuestas con GPTQ")

# Creación del formulario
with st.form("evaluation_form"):
    question = st.text_input("Pregunta", "")
    best_answer = st.text_input("Mejor Respuesta", "")
    student_answer = st.text_input("Respuesta del Estudiante", "")
    # Botón de envío para el formulario
    submitted = st.form_submit_button("Evaluar")

if submitted:
    response = generate_response(question, best_answer, student_answer)
    st.write("Respuesta del Modelo:", response)