Create app.py

app.py

@@ -0,0 +1,146 @@
+import streamlit as st
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+import json
+from pathlib import Path
+import re
+
+class WebsiteGenerator:
+    def __init__(self):
+        # Initialize the model and tokenizer
+        # Using CodeLlama as it's specifically trained for code generation
+        self.model_name = "codellama/CodeLlama-34b-Instruct-hf"
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.model_name,
+            torch_dtype=torch.float16,
+            device_map="auto"
+        )
+        
+        # Template snippets for common components
+        self.templates = {
+            "base": Path("templates/base.html").read_text(),
+            "navbar": Path("templates/navbar.html").read_text(),
+            "footer": Path("templates/footer.html").read_text()
+        }
+
+    def generate_code(self, description, component_type):
+        # Construct prompt based on component type
+        prompt = self._construct_prompt(description, component_type)
+        
+        # Generate code using the model
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+        outputs = self.model.generate(
+            inputs.input_ids,
+            max_length=2048,
+            temperature=0.7,
+            top_p=0.95,
+            do_sample=True
+        )
+        
+        generated_code = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return self._clean_generated_code(generated_code)
+
+    def _construct_prompt(self, description, component_type):
+        prompts = {
+            "frontend": f"""
+                Create a modern, responsive website frontend based on this description:
+                {description}
+                
+                Generate HTML, CSS, and JavaScript code that:
+                1. Uses modern best practices
+                2. Is mobile-responsive
+                3. Follows accessibility guidelines
+                4. Includes proper semantic HTML
+                
+                Response should include complete code:
+            """,
+            "backend": f"""
+                Create a Python FastAPI backend based on this description:
+                {description}
+                
+                Generate code that includes:
+                1. API endpoints
+                2. Database models
+                3. Authentication if needed
+                4. Basic error handling
+                
+                Response should include complete code:
+            """
+        }
+        return prompts.get(component_type, "")
+
+    def _clean_generated_code(self, code):
+        # Remove any potential unsafe or invalid code
+        code = re.sub(r'<script.*?>(.*?)</script>', '', code, flags=re.DOTALL)
+        return code
+
+def main():
+    st.set_page_config(page_title="AI Website Generator", layout="wide")
+    
+    st.title("AI Website Generator")
+    st.write("Generate complete website code using natural language descriptions")
+    
+    # Initialize generator
+    generator = WebsiteGenerator()
+    
+    # Project configuration
+    with st.expander("Project Configuration"):
+        project_name = st.text_input("Project Name")
+        framework = st.selectbox(
+            "Frontend Framework",
+            ["React", "Vue", "Plain HTML/CSS/JS"]
+        )
+        backend = st.selectbox(
+            "Backend Framework",
+            ["FastAPI", "Django", "Flask"]
+        )
+        database = st.selectbox(
+            "Database",
+            ["PostgreSQL", "MongoDB", "SQLite"]
+        )
+
+    # Website description
+    description = st.text_area(
+        "Describe your website",
+        height=200,
+        placeholder="Describe the website you want to create in detail..."
+    )
+
+    # Generate code button
+    if st.button("Generate Website Code"):
+        with st.spinner("Generating code..."):
+            # Generate frontend code
+            frontend_code = generator.generate_code(description, "frontend")
+            
+            # Generate backend code
+            backend_code = generator.generate_code(description, "backend")
+            
+            # Display generated code
+            st.subheader("Frontend Code")
+            st.code(frontend_code, language="html")
+            
+            st.subheader("Backend Code")
+            st.code(backend_code, language="python")
+            
+            # Export option
+            export_data = {
+                "project_name": project_name,
+                "frontend": frontend_code,
+                "backend": backend_code,
+                "config": {
+                    "framework": framework,
+                    "backend": backend,
+                    "database": database
+                }
+            }
+            
+            st.download_button(
+                "Download Project Files",
+                data=json.dumps(export_data, indent=2),
+                file_name=f"{project_name}_generated_code.json",
+                mime="application/json"
+            )
+
+if __name__ == "__main__":
+    main()