AstraCode / app.py
Haseeb-001's picture
Update app.py
80bea5d verified
raw
history blame
3.32 kB
import streamlit as st
from transformers import pipeline
from fpdf import FPDF # For PDF generation
# Initialize the pipeline for text generation using Llama model
llama_model = "meta-llama/Llama-2-7b-hf" # Update with the specific Llama model you want to use
# Function to get code from Hugging Face Llama model
def generate_code(summary, language):
try:
generator = pipeline('text-generation', model=llama_model, tokenizer=llama_model)
generated_code = generator(f"Generate {language} code: {summary}", max_length=150)[0]['generated_text']
return generated_code
except Exception as e:
st.error(f"Error generating code: {e}")
return ""
# Function to explain the generated code using Llama (or a suitable summarization model)
def explain_code(code):
try:
explainer = pipeline('summarization', model=llama_model, tokenizer=llama_model) # You can use a summarization model or a specific Llama-based explanation model
explanation = explainer(f"Explain the following code:\n\n{code}", max_length=250, min_length=50)[0]['summary_text']
return explanation
except Exception as e:
st.error(f"Error explaining code: {e}")
return ""
# Function to save code as a PDF
def save_code_as_pdf(code, file_name="generated_code.pdf"):
pdf = FPDF()
pdf.add_page()
pdf.set_font("Arial", size=12)
pdf.multi_cell(0, 10, code)
pdf.output(file_name)
return file_name
# --- Streamlit Interface ---
st.set_page_config(page_title="Generative AI Code Generator", page_icon="πŸ§‘β€πŸ’»", layout="wide")
# Page Title
st.title("πŸš€ Generative AI Code Generator Using Hugging Face Llama")
# Input Fields
summary = st.text_area("πŸ“ Enter the Task Summary", "For example: Create a function to add two numbers.")
language = st.selectbox("🌐 Select Programming Language", ["Python", "Java", "JavaScript", "C++"])
# Generate Code Button
if st.button("Generate Code"):
if summary:
generated_code = generate_code(summary, language)
if generated_code:
st.subheader(f"✨ Generated {language} Code:")
st.code(generated_code, language=language.lower())
# Code Modification Section
modified_code = st.text_area("✏️ Modify the Code (Optional):", value=generated_code, height=200)
# Explanation Button
if st.button("Explain Code"):
explanation = explain_code(generated_code)
st.subheader("πŸ“– Code Explanation:")
st.write(explanation)
# Download Code as PDF
if st.button("Download Code as PDF"):
pdf_path = save_code_as_pdf(modified_code) # Use modified code if edited
with open(pdf_path, "rb") as pdf_file:
st.download_button(
label="πŸ“₯ Download PDF",
data=pdf_file,
file_name="generated_code.pdf",
mime="application/pdf",
)
# New Code Button
if st.button("Generate New Code"):
st.rerun() # Refresh the page to clear inputs
# Footer Information
st.markdown("---")
st.write("🌟 Powered by **Streamlit**, **Hugging Face**, and **Transformers** | Deployed on Hugging Face")