Spaces:

shaheerawan3
/

Web_Development

Sleeping

App Files Files Community

shaheerawan3 commited on Dec 16, 2024

Commit

6294a96

verified ·

1 Parent(s): 0906714

Update app.py

Browse files

Files changed (1) hide show

app.py +92 -38

app.py CHANGED Viewed

@@ -1,46 +1,100 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-from pathlib import Path
-import re
 import json
-class WebsiteGenerator:
-    def __init__(self):
-        # Initialize the model and tokenizer
-        self.model_name = "codellama/CodeLlama-34b-Instruct-hf"
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
-        # Add memory optimization parameters
-        self.model = AutoModelForCausalLM.from_pretrained(
-            self.model_name,
-            torch_dtype=torch.float16,
-            low_cpu_mem_usage=True,
-            device_map="auto",
-            max_memory={0: "8GiB"},  # Adjust based on your GPU memory
-            offload_folder="offload"  # For memory offloading
-        )
-        # Template snippets for common components
-        self.templates = self._load_templates()
-    def _load_templates(self):
-        template_dir = Path("templates")
-        if not template_dir.exists():
-            template_dir.mkdir(exist_ok=True)
-            # Create default templates if they don't exist
-            return {
-                "base": "<!DOCTYPE html><html><head><title>Generated Site</title></head><body></body></html>",
-                "navbar": "<nav></nav>",
-                "footer": "<footer></footer>"
-            }
-        return {
-            "base": (template_dir / "base.html").read_text() if (template_dir / "base.html").exists() else "",
-            "navbar": (template_dir / "navbar.html").read_text() if (template_dir / "navbar.html").exists() else "",
-            "footer": (template_dir / "footer.html").read_text() if (template_dir / "footer.html").exists() else ""
-        }
-    # Rest of the code remains the same...
-    # (previous methods for generate_code, _construct_prompt, _clean_generated_code)
-# Main Streamlit interface code remains the same...

 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import json
+from pathlib import Path
+# Enable Streamlit caching
+@st.cache_resource
+def load_model():
+    # Using a smaller model for faster loading
+    model_name = "Salesforce/codegen-350M-mono"  # Much smaller than CodeLlama
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.float16,
+        low_cpu_mem_usage=True,
+        device_map="auto"
+    )
+    return tokenizer, model
+def generate_code(model, tokenizer, prompt, max_length=500):
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    with torch.no_grad():  # Disable gradient calculation for inference
+        outputs = model.generate(
+            inputs.input_ids,
+            max_length=max_length,
+            num_return_sequences=1,
+            temperature=0.7,
+            top_p=0.95,
+            do_sample=True
+        )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+def main():
+    st.set_page_config(page_title="Fast Web Generator", layout="wide")
+    st.title("⚡ Fast Web Generator")
+    # Initialize model with caching
+    with st.spinner("Loading model... (first run only)"):
+        tokenizer, model = load_model()
+    # Simplified interface
+    col1, col2 = st.columns([2, 1])
+    with col1:
+        description = st.text_area(
+            "Describe your website",
+            height=150,
+            placeholder="E.g., Create a responsive landing page with a navigation bar, hero section, and contact form"
+        )
+        generate_type = st.radio(
+            "What would you like to generate?",
+            ["HTML/CSS", "JavaScript", "Full Stack"],
+            horizontal=True
+        )
+    with col2:
+        st.subheader("Quick Settings")
+        use_tailwind = st.checkbox("Use Tailwind CSS", value=True)
+        responsive = st.checkbox("Make Responsive", value=True)
+        dark_mode = st.checkbox("Add Dark Mode", value=False)
+    if st.button("🚀 Generate Code"):
+        with st.spinner("Generating your code..."):
+            # Construct prompt based on settings
+            prompt_template = f"""
+            Create a website with these specifications:
+            Description: {description}
+            CSS Framework: {"Tailwind" if use_tailwind else "Plain CSS"}
+            Features: {"Responsive, " if responsive else ""}{"Dark mode, " if dark_mode else ""}
+            Generate the code:
+            """
+            generated_code = generate_code(model, tokenizer, prompt_template)
+            # Display code
+            st.code(generated_code, language="html")
+            # Add download button
+            st.download_button(
+                "📥 Download Code",
+                generated_code,
+                file_name="generated_website.html",
+                mime="text/html"
+            )
+    # Add helpful tips
+    with st.expander("💡 Tips for better results"):
+        st.markdown("""
+        - Be specific in your description
+        - Mention key features you want
+        - Specify color schemes if important
+        - Describe the layout structure
+        """)
+if __name__ == "__main__":
+    main()