import gradio as gr
from transformers import pipeline
import json

# Define the JSON structure
predefined_json = {
    "name": None,
    "age": None,
    "email": None,
    "address": None
}

# Load a lightweight model
model = pipeline("text2text-generation", model="google/flan-t5-small", max_length=256)

'''
# Load a free LLM from Hugging Face
model = pipeline("text-generation", model="tiiuae/falcon-7b-instruct", max_length=512)
'''

def convert_text_to_json(input_text):
    # Generate JSON using the LLM
    prompt = f"Extrae los campos incluidos en el siguiente formato JSON: {list(predefined_json.keys())}\nInput: {input_text}"
    response = model(prompt)[0]["generated_text"]

    # Attempt to parse the response into a JSON object
    try:
        generated_json = json.loads(response)  # Safer than eval()
    except json.JSONDecodeError:
        return {}, "Error: El modelo no retornó un formato JSON válido."

    # Check for missing fields
    missing_fields = [key for key in predefined_json if key not in generated_json or not generated_json[key]]
    missing_message = f"Campos faltantes: {', '.join(missing_fields)}" if missing_fields else "Todos los campos están completos."

    return generated_json, missing_message

# Define Gradio app
interface = gr.Interface(
    fn=convert_text_to_json,
    inputs="text",
    outputs=["json", "text"],
    title="Convertidor Texto en JSON",
    description="Ingrese el texto para extraer información en un formato JSON predefinido."
)

interface.launch()