Canstralian's picture
Update app.py
51f615c verified
import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
# Load the model and tokenizer
@st.cache_resource
def load_model():
model_name = "replit/replit-code-v1-3b" # Replace with your fine-tuned model if applicable
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
return model, tokenizer
model, tokenizer = load_model()
# App title and description
st.title("Replit-code-v1-3b Code Assistant 📊")
st.markdown("""
This application allows you to interact with the **Replit-code-v1-3b** large language model.
You can use it to generate, debug, or optimize code snippets.
Simply provide a prompt, and the model will respond with suggestions!
""")
# User input
st.header("Enter Your Prompt")
prompt = st.text_area("Describe your coding task or paste your code for debugging:")
# Temperature and max length controls
st.sidebar.header("Model Settings")
temperature = st.sidebar.slider("Temperature (Creativity)", 0.0, 1.0, 0.7)
max_length = st.sidebar.slider("Max Response Length", 50, 500, 200)
# Submit button
if st.button("Generate Response"):
if prompt.strip():
with st.spinner("Generating response..."):
# Tokenize and generate response
inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
outputs = model.generate(
inputs.input_ids,
max_length=max_length,
temperature=temperature,
pad_token_id=tokenizer.eos_token_id
)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Display response
st.subheader("Generated Code/Response")
st.code(response, language="python")
else:
st.warning("Please enter a prompt to generate a response.")
# Footer
st.markdown("---")
st.markdown("""
**Replit-code-v1-3b Code Assistant**
Built with [Streamlit](https://streamlit.io/) and the [Hugging Face Transformers Library](https://huggingface.co/docs/transformers).
""")