Spaces:
Running
Running
File size: 1,431 Bytes
7979451 43bd0e1 0f373cf 43bd0e1 |
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 |
import streamlit as st
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
# Load model and tokenizer
#model_name = "codellama/CodeLlama-7b-Python-hf"
model_name = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
st.title("Python Code Helper")
try:
st.info("Loading model... This may take a few moments.")
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
device = torch.device("cuda" if torch.cuda.is_available() else 'cpu')
model = model.to(device)
st.success("Model loaded successfully.")
except Exception as e:
st.error(f"Error loading model: {e}")
st.stop()
# Input and form handling
st.markdown("### Python Code Generation")
with st.form(key="code_form"):
prompt = st.text_area("Enter your coding prompt:", height=200)
submit = st.form_submit_button("Generate Code")
if submit and prompt.strip():
with st.spinner("Generating response..."):
try:
inputs = tokenizer(prompt, return_tensors="pt").to(device)
outputs = model.generate(**inputs, max_length=512, num_return_sequences=1)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
st.markdown("### Generated Code:")
st.code(response, language="python")
except Exception as e:
st.error(f"An error occurred: {e}")
|