import streamlit as st from transformers import AutoTokenizer, AutoModelForSeq2SeqLM @st.cache_resource def load_model(): model_name = "Salesforce/codet5-small" tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModelForSeq2SeqLM.from_pretrained(model_name) return tokenizer, model # Load the model and tokenizer (cached) with st.spinner("Loading model..."): tokenizer, model = load_model() # Streamlit UI st.title("Code Generator with Hugging Face") st.write("Generate code snippets from natural language prompts!") prompt = st.text_area("Enter your coding task:", placeholder="Write a Python function to calculate factorial.") max_length = st.slider("Select maximum length of generated code:", min_value=20, max_value=200, value=50, step=10) if st.button("Generate Code"): if prompt.strip(): with st.spinner("Generating code..."): inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True) outputs = model.generate(inputs.input_ids, max_length=max_length, num_beams=4, early_stopping=True) generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True) st.text_area("Generated Code:", generated_code, height=200) else: st.warning("Please enter a prompt!")