CodeMagic / app.py
chmawia's picture
Update app.py
9345df8 verified
import os
import streamlit as st
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import torch
import subprocess
# Force CPU usage & prevent model download issues
os.environ["HF_HOME"] = "./cache" # Store model locally
MODEL_NAME = "Salesforce/codegen-350M-mono" # Updated model
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
def generate_code(description, language):
prompt = f"Generate {language} code: {description}"
inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
outputs = model.generate(**inputs, max_length=400)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
return response.strip()
def execute_code(code, language):
if language == "Python":
try:
result = subprocess.run(['python3', '-c', code], capture_output=True, text=True, timeout=5)
return result.stdout if result.stdout else result.stderr
except Exception as e:
return str(e)
return "Code execution only supported for Python."
# Streamlit UI
st.title("Multi-Language Text-to-Code AI")
st.write("Convert natural language descriptions into code in different programming languages! Run Python code directly in the app.")
description = st.text_area("Describe your coding task...")
language = st.selectbox("Select Programming Language", ["Python", "JavaScript", "Java"])
if st.button("Generate Code"):
if description:
code = generate_code(description, language)
st.code(code, language=language.lower())
if language == "Python":
output = execute_code(code, language)
st.text_area("Execution Output", output, height=150)
else:
st.warning("Please enter a description to generate code.")