update app.py

app.py

@@ -1,3 +1,36 @@
 import gradio as gr
 
-gr.load("models/Salesforce/codet5p-220m").launch()
+gr.load("models/Salesforce/codet5p-220m").launch()
+import streamlit as st
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
+
+def main():
+    st.title("Python Code Generation App")
+
+    # Load the model and tokenizer
+    tokenizer = AutoTokenizer.from_pretrained("models/Salesforce/codet5p-220m")
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = AutoModelForSeq2SeqLM.from_pretrained("models/Salesforce/codet5p-220m").to(device)
+
+    # Get user input
+    st.subheader("Instructions")
+    st.write("Use the following format to enter prompts: Write python code for SBERT vector embedding of a sentence")
+    st.write("")    
+    query = st.text_input("Enter a prompt here: ")
+    if st.button("Generate Code"):
+        if query.strip().lower() == 'exit':
+            st.stop()
+        else:
+            # Generate summary
+            inputs = tokenizer(f"summarize:{query}", return_tensors="pt")
+            inputs = {k: v.to(device) for k, v in inputs.items()}
+            output = model.generate(**inputs, max_length=750)
+            generated_text = tokenizer.decode(output[0]).replace("summarize:", "")
+
+            # Display the generated summary
+            st.subheader("Generated Code:")
+            st.code(generated_text)
+
+if __name__ == "__main__":
+    main()