code_extraction.py

import re

#Con esta funcion podemos extraer el codigo que va a generar el modelo, no lo he probado, no se si funcione
# este codigo lo pasamos para hacer un exec() y plotearlo en Streamlit en app.py

Class Extraction():

    def extract_code(llm_output):
        
        code_match = re.search(r"```python\n(.*?)\n```", llm_output, re.DOTALL)
        if code_match:
            return code_match.group(1)
        return None