Update app.py

app.py

@@ -1,140 +1,64 @@
-
-import os
 import streamlit as st
-from groq import Groq
-
-# Set the API key
-GROQ_API_KEY = "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"
-
-# Initialize the Groq client
-client = Groq(api_key=GROQ_API_KEY)
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import groq
 
-# Initialize session state for conversation history
-if 'conversation' not in st.session_state:
-    st.session_state.conversation = []
+# Constants
+PRIMARY_MODEL = "Qwen/CodeQwen1.5-7B-Chat"
+BACKUP_MODEL = "llama-3.3-70b-versatile"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Streamlit UI setup
-st.set_page_config(page_title="Code Generator", layout="wide")
-st.title("Code Generator")
+# Load Primary Model & Tokenizer
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained(PRIMARY_MODEL)
+    model = AutoModelForCausalLM.from_pretrained(PRIMARY_MODEL, device_map="auto")
+    return model, tokenizer
 
-# Custom CSS for neon theme
-st.markdown(
-    """
-    <style>
-    body {
-        background-color: #000;
-        color: #fff;
-        font-family: 'Courier New', Courier, monospace;
-    }
-    .stButton > button {
-        background-color: #5e72e4;
-        color: white;
-        border: none;
-        padding: 10px 20px;
-        border-radius: 5px;
-        cursor: pointer;
-        transition: background-color 0.3s ease;
-    }
-    .stButton > button:hover {
-        background-color: #324ab2;
-    }
-    .stTextInput > div > div > input {
-        background-color: #1a1a1a;
-        color: #fff;
-        border: 1px solid #5e72e4;
-        border-radius: 5px;
-        padding: 10px;
-    }
-    .stMarkdown {
-        color: #5e72e4;
-    }
-    </style>
-    """,
-    unsafe_allow_html=True
-)
+model, tokenizer = load_model()
 
-# Function to generate code
-def generate_code(query):
-    messages = [
-        {"role": "user", "content": query}
-    ]
-    messages = st.session_state.conversation + messages
-    completion = client.chat.completions.create(
-        model="qwen-2.5-coder-32b",
-        messages=messages,
-        temperature=0.6,
-        max_completion_tokens=4096,
-        top_p=0.95,
-        stream=False,
-        stop=None,
+def generate_code(prompt):
+    """Generate code from prompt using the primary model."""
+    try:
+        inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
+        outputs = model.generate(**inputs, max_length=512)
+        return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    except Exception as e:
+        st.error("Primary model failed. Switching to backup...")
+        return generate_code_backup(prompt)
+
+def generate_code_backup(prompt):
+    """Use Groq API for backup model code generation."""
+    client = groq.Client(api_key="YOUR_GROQ_API_KEY")
+    response = client.chat.completions.create(
+        model=BACKUP_MODEL,
+        messages=[{"role": "system", "content": "You are a helpful AI."},
+                  {"role": "user", "content": prompt}]
     )
-    return completion.choices[0].message.content
+    return response.choices[0].message.content
 
-# Function to check code correctness
-def check_code_correctness(code):
-    # Placeholder for code correctness check
-    # This can be replaced with an actual code checking mechanism
-    return True
-
-# Function to explain code
 def explain_code(code):
-    messages = [
-        {"role": "user", "content": f"Explain the following code:\n{code}"}
-    ]
-    completion = client.chat.completions.create(
-        model="llama-3.3-70b-versatile",
-        messages=messages,
-        temperature=0.6,
-        max_completion_tokens=4096,
-        top_p=0.95,
-        stream=False,
-        stop=None,
-    )
-    return completion.choices[0].message.content
+    """Generate a detailed explanation of the given code."""
+    prompt = f"Explain this code in detail:\n{code}"
+    return generate_code(prompt)
 
-# Main UI components
-query = st.text_area("Enter your query/requirement/demand for the app:", height=100)
+# Streamlit UI
+st.set_page_config(page_title="AI Code Generator", layout="wide")
+st.title("🚀 AI Code Generator App")
+st.markdown("### Generate and edit code with AI!")
 
-if st.button("Generate Code"):
-    try:
-        generated_code = generate_code(query)
-        st.session_state.conversation.append({"role": "assistant", "content": generated_code})
-        st.code(generated_code, language="python")
-
-        # Placeholder for multiple solutions
-        # You can replace this with actual logic to fetch multiple solutions
-        alternative_solutions = ["Alternative Solution 1", "Alternative Solution 2"]
-        if alternative_solutions:
-            st.sidebar.header("Alternative Solutions")
-            for idx, solution in enumerate(alternative_solutions, start=1):
-                st.sidebar.write(f"**Solution {idx}:** {solution}")
+# Sidebar for Alternative Solutions
+st.sidebar.header("⚡ Alternative Solutions")
 
-        # Check code correctness
-        if check_code_correctness(generated_code):
-            st.success("The generated code is correct.")
-        else:
-            st.error("The generated code may contain errors.")
-
-    except Exception as e:
-        st.error(f"An error occurred: {e}")
+user_query = st.text_area("Enter your app idea:", height=150)
 
-if st.button("Explain Code"):
-    try:
-        code_to_explain = st.text_area("Enter the code you want to explain:", height=200)
-        explanation = explain_code(code_to_explain)
-        st.markdown(f"### Explanation:\n{explanation}")
-    except Exception as e:
-        st.error(f"An error occurred: {e}")
+if st.button("Generate Code"):
+    with st.spinner("Generating code..."):
+        generated_code = generate_code(user_query)
+        st.code(generated_code, language="python")
+        st.session_state["generated_code"] = generated_code
 
-if st.button("New Code"):
-    st.session_state.conversation = []
-    st.experimental_rerun()
+if "generated_code" in st.session_state and st.button("📜 Explain Code"):
+    explanation = explain_code(st.session_state["generated_code"])
+    st.text_area("Explanation:", explanation, height=250)
 
-# Display conversation history
-if st.session_state.conversation:
-    st.sidebar.header("Conversation History")
-    for message in st.session_state.conversation:
-        if message["role"] == "user":
-            st.sidebar.markdown(f"**User:** {message['content']}")
-        else:
-            st.sidebar.markdown(f"**Assistant:** {message['content']}")
+st.sidebar.button("🔄 New Code", on_click=lambda: st.experimental_rerun())