File size: 4,340 Bytes
5dd1fa5 a675706 5dd1fa5 a675706 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | import streamlit as st
import pinecone
import os
import subprocess
import json
import requests
from pinecone import Pinecone, ServerlessSpec
#from streamlit_extras.add_vertical_space import add_vertical_space
# Set page config
st.set_page_config(page_title="StarCoder AI Code Generator", page_icon="🤖", layout="wide")
# Custom CSS for Grok-like UI
st.markdown("""
<style>
body {background-color: #121212; color: white;}
.stApp {background: #181818; color: white; font-family: 'Arial', sans-serif;}
.stTextArea textarea {background: #282828; color: white; border-radius: 10px;}
.stTextInput input {background: #282828; color: white; border-radius: 10px;}
.stButton>button {background: #3f51b5; color: white; border-radius: 8px; padding: 10px;}
.stDownloadButton>button {background: #4caf50; color: white; border-radius: 8px; padding: 10px;}
.stCode {background: #202020; padding: 15px; border-radius: 10px;}
.stSidebar {background: #202020; color: white;}
.stSidebar .stButton>button {background: #ff9800; color: white;}
</style>
""", unsafe_allow_html=True)
# Pinecone API Key
PINECONE_API_KEY = "your pinecone api key here"
pc = Pinecone(api_key=PINECONE_API_KEY)
INDEX_NAME = "code-gen"
if INDEX_NAME not in pc.list_indexes().names():
pc.create_index(name=INDEX_NAME, dimension=1536, metric='euclidean',
spec=ServerlessSpec(cloud='aws', region='us-east-1'))
index = pc.Index(INDEX_NAME)
# Hugging Face API Key
HUGGINGFACE_API_KEY = "your huggingface api key"
# Initialize session state
if "chat_history" not in st.session_state:
st.session_state.chat_history = []
if "current_chat" not in st.session_state:
st.session_state.current_chat = []
def new_chat():
if st.session_state.current_chat:
st.session_state.chat_history.append(st.session_state.current_chat)
st.session_state.current_chat = []
# Sidebar - Chat History
st.sidebar.title("💬 Chat History")
for i, chat in enumerate(st.session_state.chat_history):
if st.sidebar.button(f"Chat {i + 1} 🕒"):
st.session_state.current_chat = chat
# Sidebar - New Chat Button
if st.sidebar.button("New Chat ➕"):
new_chat()
# Function to query Hugging Face model
def generate_code(prompt, examples=[]):
formatted_prompt = "".join(examples) + "\n" + prompt
headers = {"Authorization": f"Bearer {HUGGINGFACE_API_KEY}"}
data = {"inputs": formatted_prompt, "parameters": {"temperature": 0.5, "max_length": 512}}
response = requests.post("https://api-inference.huggingface.co/models/bigcode/starcoder2-15b", headers=headers,
json=data)
return response.json()[0][
'generated_text'].strip() if response.status_code == 200 else "Error: Unable to generate code"
def execute_python_code(code):
try:
result = subprocess.run(["python", "-c", code], capture_output=True, text=True, timeout=5)
return result.stdout if result.returncode == 0 else result.stderr
except Exception as e:
return str(e)
def store_in_pinecone(prompt, code):
vector = [0.1] * 1536 # Placeholder vector
index.upsert([(prompt, vector, {"code": code})])
def retrieve_from_pinecone(prompt):
response = index.query(vector=[0.1] * 1536, top_k=2, include_metadata=True)
return [r["metadata"]["code"] for r in response["matches"]]
# Main UI
st.title("🤖 StarCoder AI Code Generator")
# User Input
prompt = st.text_area("✍️ Enter your prompt:", height=100)
mode = st.selectbox("⚡ Choose prompting mode", ["One-shot", "Two-shot"])
generate_button = st.button("🚀 Generate Code")
if generate_button and prompt:
examples = retrieve_from_pinecone(prompt) if mode == "Two-shot" else []
generated_code = generate_code(prompt, examples)
st.session_state.current_chat.append((prompt, generated_code))
st.subheader("📜 Generated Code")
st.code(generated_code, language="python")
store_in_pinecone(prompt, generated_code)
execute_button = st.button("▶️ Run Code")
if execute_button:
execution_result = execute_python_code(generated_code)
st.subheader("🖥️ Execution Output")
st.text(execution_result)
st.download_button("📥 Download Code", generated_code, file_name="generated_code.py")
|