Update app.py

app.py

@@ -1,27 +1,23 @@
 import os
 import streamlit as st
-import transformers
-from transformers import AutoModelForCausalLM
-
-os.system("pip freeze > requirements.txt")
-os.system("pip install -r requirements.txt")
+import torch
+from transformers import GPTNeoForCausalLM, GPT2Tokenizer
 
 # Function to reverse engineer code
 def reverse_prompt_engineer(input_code):
-    # Analyze the input code using langchain
-    analyzed_code = input_code
-
     # Load the tokenizer and model
-    tokenizer = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B")
-    model = tokenizer.model
+    tokenizer = GPT2Tokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")
+    model = GPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B").to("cuda")
     
     # Generate a prompt using the analyzed code
-    prompt = "Reverse engineer the following code:\n\n" + analyzed_code
+    prompt = "Reverse engineer the following code:\n\n" + input_code
     
-    # Generate similar code using the ChatGPT agent's generate_code method
-    generated_code = tokenizer(prompt, max_length=100, do_sample=True)
+    # Tokenize the prompt and generate similar code using the model
+    input_ids = tokenizer.encode(prompt, return_tensors="pt").to("cuda")
+    generated_ids = model.generate(input_ids, max_length=100, do_sample=True)
+    generated_code = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
     
-    return generated_code[0]['generated_text']
+    return generated_code
 
 # Set Streamlit page configuration
 st.set_page_config(