Spaces:
Sleeping
Sleeping
File size: 1,649 Bytes
7979451 43bd0e1 d956f4b 0f373cf 43bd0e1 d956f4b 1fff4b1 8d3a00d 367d8e0 1fff4b1 43bd0e1 d956f4b 43bd0e1 d956f4b 43bd0e1 d956f4b |
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 38 39 40 41 42 43 44 |
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:
info = st.empty()
info.markdown("#### :red[Model is Loading....]")
model = AutoModelForCausalLM.from_pretrained(model_name, revision="main")
tokenizer = AutoTokenizer.from_pretrained(model_name, revision="main", use_fast=True)
# 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)
info.markdown("#### :green[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}") |