Update app.py

app.py

@@ -1,22 +1,29 @@
-import requests
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-# Replace with the correct Space API URL
-api_url = "https://hf.space/embed/Canstralian/WhiteRabbitNeo/api/predict/"
+# Load the model and tokenizer
+model_path = "Canstralian/pentest_ai"
+model = AutoModelForCausalLM.from_pretrained(model_path)
+tokenizer = AutoTokenizer.from_pretrained(model_path)
 
-# The input payload to send to the model
-payload = {"data": ["A new bug bounty to solve"]}
-
-try:
-    # Make a POST request to the API
-    response = requests.post(api_url, json=payload)
-    response.raise_for_status()  # Raise an exception for HTTP errors
+# Function to handle user inputs and generate responses
+def generate_text(instruction):
+    # Encode the input and ensure it fits within model limits
+    inputs = tokenizer.encode(instruction, return_tensors='pt', truncation=True, max_length=512)
     
-    # Parse the result
-    result = response.json()  # Parse the JSON response from the API
-    print("Generated Result:", result["data"][0])  # Access the generated result
+    # Generate the output with additional parameters for better control
+    outputs = model.generate(inputs, max_length=150, num_beams=5, temperature=0.7, top_p=0.95, do_sample=True)
     
-    # Optionally, save the result to a text file
-    with open("generated_output.txt", "w") as f:
-        f.write(result["data"][0])
-except requests.exceptions.RequestException as e:
-    print("An error occurred:", e)
+    # Decode and return the result
+    output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return output_text
+
+# Create a Gradio interface
+iface = gr.Interface(fn=generate_text, 
+                     inputs=gr.Textbox(lines=2, placeholder="Enter your question here..."), 
+                     outputs="text", 
+                     live=True)
+
+# Launch the interface
+iface.launch()