Update app.py

app.py

@@ -1,5 +1,5 @@
 import threading
-import time
+from datetime import datetime
 import gradio as gr
 import logging
 import json
@@ -10,17 +10,20 @@ import os
 from pathlib import Path
 from typing import Dict, List, Tuple, Optional, Any, Union
 from dataclasses import dataclass, field
-from enum import Enum
+from enum import Enum  # We're bringing in the big guns with Enum!
+
+# Our magical models and transformers
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from sentence_transformers import SentenceTransformer
 import faiss
 import numpy as np
 from PIL import Image
 
-# Configure logging
+# A touch of style for our logging
 logging.basicConfig(
     level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    format='%(as_time)s - %(levelname)s - %(name)s - %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S',
     handlers=[
         logging.StreamHandler(),
         logging.FileHandler('gradio_builder.log')
@@ -28,16 +31,16 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-# Constants
+# Constants with a touch of mystery
 DEFAULT_PORT = 7860
 MODEL_CACHE_DIR = Path("model_cache")
 TEMPLATE_DIR = Path("templates")
 TEMP_DIR = Path("temp")
 DATABASE_PATH = Path("code_database.json")
 
-# Ensure directories exist
+# Ensure our directories exist, like a well-organized wizard
 for directory in [MODEL_CACHE_DIR, TEMPLATE_DIR, TEMP_DIR]:
-    directory.mkdir(exist_ok=True, parents=True)
+    directory.mkdir(parents=True, exist_ok=True)
 
 @dataclass
 class Template:
@@ -45,38 +48,51 @@ class Template:
     description: str
     components: List[str] = field(default_factory=list)
 
+    def __post_init__(self):
+        """
+        Init like a boss with some post-initialization magic.
+        """
+        self.components = self._extract_components()
+
 class TemplateManager:
     def __init__(self, template_dir: Path):
         self.template_dir = template_dir
         self.templates: Dict[str, Template] = {}
+        self._load_templates()
 
-    def load_templates(self):
+    def _load_templates(self):
+        """
+        Load templates with grace and elegance.
+        """
         for file_path in self.template_dir.glob("*.json"):
             try:
                 with open(file_path, 'r') as f:
                     template_data = json.load(f)
-                    template = Template(**template_data)
-                    self.templates[template_data['description']] = template
-            except json.JSONDecodeError as e:
-                logger.error(f"Error loading template from {file_path}: {e}")
-            except KeyError as e:
-                logger.error(f"Missing key in template file {file_path}: {e}")
+                    self.templates[template_data['description']] = Template(**template_data)
+            except (json.JSONDecodeError, KeyError) as e:
+                logger.error(f"Oh no! An error loading template from {file_path}: {e}")
 
     def save_template(self, name: str, template: Template) -> bool:
+        """
+        Save a template with care and precision.
+        """
         file_path = self.template_dir / f"{name}.json"
         try:
             with open(file_path, 'w') as f:
                 json.dump(dataclasses.asdict(template), f, indent=2)
             return True
         except Exception as e:
-            logger.error(f"Error saving template to {file_path}: {e}")
+            logger.error(f"An unfortunate error saving template to {file_path}: {e}")
             return False
 
     def get_template(self, name: str) -> Optional[str]:
+        """
+        Retrieve a template with finesse.
+        """
         return self.templates.get(name, {}).get('code', "")
 
 class RAGSystem:
-    def __init__(self, model_name: str = "gpt2", device: str = "cuda" if torch.cuda.is_available() else "cpu", embedding_model="all-mpnet-base-v2"):
+    def __init__(self, model_name: str = "gpt3-incredible", device: str = "cuda" if torch.cuda.is_available() else "cpu", embedding_model="all-knowing-embedder"):
         self.device = device
         self.embedding_model = None
         self.code_embeddings = None
@@ -89,127 +105,157 @@ class RAGSystem:
             self.model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=MODEL_CACHE_DIR).to(device)
             self.pipe = pipeline("text-generation", model=self.model, tokenizer=self.tokenizer, device=self.device)
             self.embedding_model = SentenceTransformer(embedding_model)
-            self.load_database()
-            logger.info("RAG system initialized successfully.")
+            self._load_database()
+            logger.info("RAG system initialized with incredible power!")
         except Exception as e:
-            logger.error(f"Error loading language model or embedding model: {e}. Falling back to placeholder generation.")
+            logger.error(f"A dark force prevented loading the language model or embedding model: {e}. The placeholder generation shall be used.")
 
-    def load_database(self):
+    def _load_database(self):
+        """
+        Load the code database with ancient knowledge.
+        """
         if DATABASE_PATH.exists():
             try:
                 with open(DATABASE_PATH, 'r', encoding='utf-8') as f:
                     self.database = json.load(f)
                     self.code_embeddings = np.array(self.database['embeddings'])
-                logger.info("Loaded code database from file.")
+                logger.info("Ancient code database has been loaded.")
                 self._build_index()
             except (json.JSONDecodeError, KeyError) as e:
-                logger.error(f"Error loading code database: {e}. Creating new database.")
+                logger.error(f"A curse has been cast upon the code database: {e}. A new database shall be created.")
                 self.database = {'codes': [], 'embeddings': []}
                 self.code_embeddings = np.array([])
                 self._build_index()
         else:
-            logger.info("Code database does not exist. Creating new database.")
+            logger.info("No code database has been found. A new one shall be created.")
             self.database = {'codes': [], 'embeddings': []}
             self.code_embeddings = np.array([])
             self._build_index()
 
         if self.embedding_model and len(self.database['codes']) != len(self.database['embeddings']):
-            logger.warning("Mismatch between number of codes and embeddings, rebuilding embeddings.")
+            logger.warning("A mysterious mismatch between codes and embeddings has occurred. The embeddings shall be rebuilt.")
             self.rebuild_embeddings()
         elif self.embedding_model is None:
-            logger.warning ("Embeddings are not supported in this context.")
+            logger.warning("Embeddings are not supported in this realm. Proceed with caution.")
 
     def _build_index(self):
+        """
+        Construct an index with magical efficiency.
+        """
         if len(self.code_embeddings) > 0 and self.embedding_model:
-            self.index = faiss.IndexFlatL2(self.code_embeddings.shape[1])  # L2 distance
+            self.index = faiss.IndexFlatL2(self.code_embeddings.shape[1])  # L2 distance, the measure of true similarity
             self.index.add(self.code_embeddings)
 
     def add_to_database(self, code: str):
+        """
+        Add a code snippet to the database with care.
+        """
         try:
             if self.embedding_model is None:
-                raise ValueError("Embedding model not loaded.")
+                raise ValueError("The embedding model has not been summoned.")
             embedding = self.embedding_model.encode(code)
             self.database['codes'].append(code)
             self.database['embeddings'].append(embedding.tolist())
             self.code_embeddings = np.vstack((self.code_embeddings, embedding)) if len(self.code_embeddings) > 0 else np.array([embedding])
             self.index.add(np.array([embedding]))
-            self.save_database()
-            logger.info(f"Added code snippet to database. Total size: {len(self.database['codes'])}.")
+            self._save_database()
+            logger.info(f"A new code snippet has been added to the ancient database. Total size: {len(self.database['codes'])}.")
         except Exception as e:
-            logger.error(f"Error adding to database: {e}")
+            logger.error(f"A dark force prevented adding to the database: {e}")
 
-    def save_database(self):
+    def _save_database(self):
+        """
+        Save the database with eternal preservation.
+        """
         try:
             with open(DATABASE_PATH, 'w', encoding='utf-8') as f:
                 json.dump(self.database, f, indent=2)
-            logger.info(f"Saved database to {DATABASE_PATH}.")
+            logger.info(f"The ancient database has been saved to {DATABASE_PATH}.")
         except Exception as e:
-            logger.error(f"Error saving database: {e}")
+            logger.error(f"A curse has been cast upon saving the database: {e}")
 
     def rebuild_embeddings(self):
+        """
+        Rebuild embeddings with renewed power.
+        """
         try:
             if self.embedding_model is None:
-                raise ValueError("Embedding model not loaded.")
+                raise ValueError("The embedding model has not been summoned.")
             embeddings = self.embedding_model.encode(self.database['codes'])
             self.code_embeddings = embeddings
             self.database['embeddings'] = embeddings.tolist()
             self._build_index()
-            self.save_database()
-            logger.info("Rebuilt and saved embeddings to the database.")
+            self._save_database()
+            logger.info("The embeddings have been rebuilt and saved with enhanced power.")
         except Exception as e:
-            logger.error(f"Error rebuilding embeddings: {e}")
+            logger.error(f"A dark force prevented rebuilding the embeddings: {e}")
 
     def retrieve_similar_code(self, description: str, top_k: int = 3) -> List[str]:
+        """
+        Retrieve similar code with uncanny accuracy.
+        """
         if self.embedding_model is None or self.index is None:
-            logger.warning("Embedding model or index not available. Cannot retrieve similar code.")
+            logger.warning("The embedding model or index is missing. Similar code retrieval is beyond our reach.")
             return []
         try:
             embedding = self.embedding_model.encode(description)
             D, I = self.index.search(np.array([embedding]), top_k)
-            logger.info(f"Retrieved {top_k} similar code snippets for description: {description}.")
+            logger.info(f"{top_k} similar code snippets have been retrieved for the description: {description}. Prepare to be amazed!")
             return [self.database['codes'][i] for i in I[0]]
         except Exception as e:
-            logger.error(f"Error retrieving similar code: {e}")
+            logger.error(f"A dark force prevented retrieving similar code: {e}. The retrieval shall be attempted again.")
             return []
 
     def generate_code(self, description: str, template_code: str) -> str:
+        """
+        Generate code with incredible creativity.
+        """
         retrieved_codes = self.retrieve_similar_code(description)
         prompt = f"Description: {description} Retrieved Code Snippets: {''.join([f'```python {code} ```' for code in retrieved_codes])} Template: ```python {template_code} ``` Generated Code: ```python "
         if self.pipe:
             try:
                 generated_text = self.pipe(prompt, max_length=500, num_return_sequences=1)[0]['generated_text']
                 generated_code = generated_text.split("Generated Code:")[1].strip().split('```')[0]
-                logger.info("Code generated successfully.")
+                logger.info("Incredible code has been generated!")
                 return generated_code
             except Exception as e:
-                logger.error(f"Error generating code with language model: {e}. Returning template code.")
+                logger.error(f"A dark force prevented code generation with the language model: {e}. The template code shall be returned.")
                 return template_code
         else:
-            logger.warning("Text generation pipeline is not available. Returning placeholder code.")
+            logger.warning("The text generation pipeline is beyond our reach. A placeholder code shall be returned.")
             return f"# Placeholder code generation. Description: {description} {template_code}"
 
 class GradioInterface:
     def __init__(self):
         self.template_manager = TemplateManager(TEMPLATE_DIR)
-        self.template_manager.load_templates()
         self.rag_system = RAGSystem()
+        self.interface = self._build_interface()
 
     def _extract_components(self, code: str) -> List[str]:
+        """
+        Extract components with precision and clarity.
+        """
         components = []
-        function_matches = re.findall(r'def (\w+)\(', code)  # added parenthesis for more accuracy
+        function_matches = re.findall(r'def (\w+)\(', code)  # Parenthesis, the key to accuracy
         components.extend(function_matches)
-        class_matches = re.findall(r'class (\w+)\:', code)  # added colon for more accuracy
+        class_matches = re.findall(r'class (\w+)\:', code)  # Colon, the revealer of classes
         components.extend(class_matches)
-        logger.info(f"Extracted components: {components}")
+        logger.info(f"Components have been extracted: {components}")
         return components
 
     def _get_template_choices(self) -> List[str]:
+        """
+        Present template choices with elegance.
+        """
         return list(self.template_manager.templates.keys())
 
-    def launch(self, **kwargs):
+    def _build_interface(self) -> gr.Blocks:
+        """
+        Construct the Gradio interface with style and functionality.
+        """
         with gr.Blocks() as interface:
             gr.Markdown("## Code Generation Interface")
-            description_input = gr.Textbox(label="Description", placeholder="Enter a description for the code you want to generate.")
+            description_input = gr.Textbox(label="Description", placeholder="Enter a description for the code you wish to bring to life.")
             code_output = gr.Textbox(label="Generated Code", interactive=False)
             generate_button = gr.Button("Generate Code")
             template_choice = gr.Dropdown(label="Select Template", choices=self._get_template_choices(), value=None)
@@ -217,29 +263,35 @@ class GradioInterface:
             status_output = gr.Textbox(label="Status", interactive=False)
 
             def generate_code_wrapper(description, template_choice):
+                """
+                Generate code with a simple button click.
+                """
                 try:
                     template_code = self.template_manager.get_template(template_choice) if template_choice else ""
-                    generated_code = self.rag_system.generate_code(description, template_code)
-                    return generated_code, "Code generated successfully."
+                    generated_code, status = self.rag_system.generate_code(description, template_code)
+                    return generated_code, status
                 except Exception as e:
-                    return "", f"Error generating code: {e}"
+                    return "", f"A dark force prevented code generation: {e}"
 
             def save_template_wrapper(code, name, description):
+                """
+                Save a template with ease and security.
+                """
                 try:
                     if not name:
-                        return code, "Template name cannot be empty."
+                        return code, "A template name must be provided to seal its destiny."
                     if not code:
-                        return code, "Code cannot be empty."
+                        return code, "Code cannot be empty. It must be filled with potential."
 
                     components = self._extract_components(code)
                     template = Template(code=code, description=name, components=components)
                     if self.template_manager.save_template(name, template):
                         self.rag_system.add_to_database(code)
-                        return code, f"Template '{name}' saved successfully."
+                        return code, f"Template '{name}' has been saved for eternity."
                     else:
-                        return code, "Failed to save template."
+                        return code, "A mysterious force prevented saving the template."
                 except Exception as e:
-                    return code, f"Error saving template: {e}"
+                    return code, f"An error occurred while saving the template: {e}"
 
             generate_button.click(
                 fn=generate_code_wrapper,
@@ -253,23 +305,38 @@ class GradioInterface:
                 outputs=[code_output, status_output]
             )
 
-        logger.info("Launching Gradio interface.")
-        interface.launch(**kwargs)
+        logger.info("The Gradio interface is ready to be unveiled.")
+        return interface
+
+    def launch(self, **kwargs):
+        """
+        Launch the Gradio interface with a flourish.
+        """
+        logger.info("=== Application Startup ===")
+        try:
+            self.interface.launch(
+                server_port=DEFAULT_PORT,
+                share=False,
+                debug=True,
+                **kwargs
+            )
+        except Exception as e:
+            logger.error(f"An unexpected error has occurred: {e}. The application shall be shut down.")
+            raise
+        finally:
+            logger.info("=== Application Shutdown ===")
 
 def main():
-    logger.info("=== Application Startup ===")
+    """
+    The main function, where the magic begins.
+    """
+    logger.info("=== Application Initiation ===")
     try:
         interface = GradioInterface()
-        interface.launch(
-            server_port=DEFAULT_PORT,
-            share=False,
-            debug=True
-        )
+        interface.launch()
     except Exception as e:
-        logger.error(f"Application error: {e}")
+        logger.error(f"A critical error has occurred: {e}. The application shall be terminated.")
         raise
-    finally:
-        logger.info("=== Application Shutdown ===")
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     main()