Update app.py

app.py

@@ -1,12 +1,18 @@
+import os
 import streamlit as st
-from transformers import T5ForConditionalGeneration, T5Tokenizer  # Replace with the library used for your code generation model
-import sentencepiece as spm
-# Load your code generation model
-model_path = "Cegil/code_generation"  # e.g., a model checkpoint or saved weights
+from transformers import T5ForConditionalGeneration, T5Tokenizer
+import torch
+
+# Define the model path and check if it exists
+model_path = "Cegil/code_generation"  # Update to your model's actual path
+if not os.path.exists(model_path):
+    raise FileNotFoundError(f"Model path '{model_path}' does not exist.")
+
+# Load the model and tokenizer
 code_gen_model = T5ForConditionalGeneration.from_pretrained(model_path)
 tokenizer = T5Tokenizer.from_pretrained(model_path)
 
-# UI for Code Generation
+# Streamlit UI
 st.title("Code Generation Interface")
 
 # Input prompt for code generation
@@ -14,9 +20,15 @@ prompt = st.text_input("Enter your code generation prompt:", "Example prompt")
 
 # Button to generate code
 if st.button("Generate Code"):
-    # Use the code generation model to generate code based on the prompt
-    generated_code = code_gen_model.generate(prompt)  # Adapt this to your model's method
+    # Tokenize the input prompt
+    inputs = tokenizer(prompt, return_tensors="pt")  # Ensure input is properly tokenized
+
+    # Generate code using the code generation model
+    output = code_gen_model.generate(inputs['input_ids'])
+
+    # Decode the output to get the generated code
+    generated_code = tokenizer.decode(output[0], skip_special_tokens=True)
 
-    # Display the generated code
+    # Display the generated code in a formatted way
     st.write("Generated Code:")
     st.code(generated_code)