commit

app.py

@@ -2,46 +2,34 @@ import streamlit as st
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-
+# Load model and tokenizer
 model_name = "codellama/CodeLlama-7b-Python-hf"
-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)
-
-# prompt = st.text_area("Enter your prompt:")
-
-
-# def translate_text(text, source_lang, target_lang):
-#     tokenizer.src_lang = source_lang
-#     encoded_text = tokenizer(text, return_tensors="pt").to(device)
-
-#     generated_tokens = model.generate(**encoded_text, forced_bos_token_id=tokenizer.lang_code_to_id[target_lang])
-
-#     #Decode the output
-#     translated_text = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
-#     return translated_text
-
-
-st.markdown("### Python Code Helper")
-# source_language = ''
-# target_language = ''
-# source = st.sidebar.selectbox('Source Language', languages)
-# if source:
-#     source_language = lang_dict.get(source)
-#     st.write(source_language)
-
-# target = st.sidebar.selectbox('Target Language', languages)
-# if target:
-#     target_language = lang_dict.get(target)
-#     st.write(target_language)
-
-with st.form(key="myForm"):
-    prompt = st.text_area("Enter your Prompt")
-    submit = st.form_submit_button("Submit", type='primary')
-    if submit and prompt:
-        with st.spinner("Generating Response"):
-            response = model.invoke(prompt)
-            st.write(response)
-        
+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}")