Update app.py

app.py

@@ -1,26 +1,31 @@
-import streamlit as st
-from transformers import pipeline
+import gradio as gr
+from huggingface_hub import InferenceClient
 
-# Загружаем модель для генерации текста
-pipe = pipeline("text-generation", model="WhiteRabbitNeo/WhiteRabbitNeo-13B-v1", trust_remote_code=True)
+# Идентификатор модели
+model_id = "WhiteRabbitNeo/WhiteRabbitNeo-13B-v1"
+
+
+
+# Создание клиента для доступа к модели
+client = InferenceClient(model_id, token=HF_TOKEN)
 
 # Основная функция для разведки
 def reconnaissance(domain):
     st.write(f"Начинаем разведку для домена: {domain}")
-    result = pipe(f"Разведка информации о домене {domain}...")
-    st.text_area("Результаты разведки", result[0]['generated_text'])
+    result = client.predict(f"Разведка информации о домене {domain}...")
+    st.text_area("Результаты разведки", result <sup> </sup>['generated_text'])
 
 # Функция для поиска уязвимостей
 def vulnerability_search(service_info):
     st.write(f"Ищем уязвимости для сервиса: {service_info}")
-    result = pipe(f"Поиск уязвимостей для сервиса {service_info}...")
-    st.text_area("Результаты поиска уязвимостей", result[0]['generated_text'])
+    result = client.predict(f"Поиск уязвимостей для сервиса {service_info}...")
+    st.text_area("Результаты поиска уязвимостей", result <sup> </sup>['generated_text'])
 
 # Функция для создания эксплойта
 def exploit_creation(vulnerability, language):
     st.write(f"Создание эксплойта для уязвимости: {vulnerability} на языке {language}")
-    result = pipe(f"Генерация эксплойта для уязвимости {vulnerability} на языке {language}...")
-    st.text_area("Генерированный эксплойт", result[0]['generated_text'])
+    result = client.predict(f"Генерация эксплойта для уязвимости {vulnerability} на языке {language}...")
+    st.text_area("Генерированный эксплойт", result <sup> </sup>['generated_text'])
 
 # Интерфейс Streamlit
 st.title("Инструмент для Red Team операций с WhiteRabbitNeo")
@@ -44,3 +49,7 @@ elif task == "Создание эксплойта":
     language = st.selectbox("Выберите язык для эксплойта:", ["Python", "Bash", "PowerShell"])
     if st.button("Создать эксплойт"):
         exploit_creation(vulnerability, language)
+
+# Запуск приложения
+if __name__ == "__main__":
+    demo.launch()