Update app.py

app.py

@@ -1,60 +1,53 @@
-import spaces
 import gradio as gr
-import torch
-import transformers
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-import os
-
-title = """# Welcome to 🌟Tonic's✨StarCoder
-✨StarCoder StarCoder2-15B model is a 15B parameter model trained on 600+ programming languages from The Stack v2, with opt-out requests excluded. The model uses Grouped Query Attention, a context window of 16,384 tokens with a sliding window attention of 4,096 tokens, and was trained using the Fill-in-the-Middle objective on 4+ trillion tokens. The model was trained with NVIDIA NeMo™ Framework using the NVIDIA Eos Supercomputer built with NVIDIA DGX H100 systems. You can build with this endpoint using✨StarCoder available here : [bigcode/starcoder2-15b](https://huggingface.co/bigcode/starcoder2-15b). You can also use ✨StarCoder by cloning this space. Simply click here: <a style="display:inline-block" href="https://huggingface.co/spaces/Tonic/starcoder2?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></h3> 
-Join us : 🌟TeamTonic🌟 is always making cool demos! Join our active builder's 🛠️community 👻 [![Join us on Discord](https://img.shields.io/discord/1109943800132010065?label=Discord&logo=discord&style=flat-square)](https://discord.gg/GWpVpekp) On 🤗Huggingface:[MultiTransformer](https://huggingface.co/MultiTransformer) Math 🔍 [introspector](https://huggingface.co/introspector) On 🌐Github: [Tonic-AI](https://github.com/tonic-ai) & contribute to🌟 [SciTonic](https://github.com/Tonic-AI/scitonic)🤗Big thanks to Yuvi Sharma and all the folks at huggingface for the community grant 🤗
-"""
-
-model_path = "bigcode/starcoder2-15b"
-
-hf_token = os.getenv("HF_TOKEN")
-if not hf_token:
-    raise ValueError("Hugging Face token not found. Please set the HF_TOKEN environment variable.")
-
-tokenizer = AutoTokenizer.from_pretrained(model_path)
-quantization_config = BitsAndBytesConfig(load_in_8bit=True)
-model = AutoModelForCausalLM.from_pretrained( model_path, quantization_config=quantization_config)
-
-@spaces.GPU
-def generate_text(prompt, temperature=0.9, max_length=1200):
-    # Encode the inputs
-    inputs = tokenizer.encode(prompt, return_tensors="pt")
-    attention_mask = torch.ones(inputs.shape, dtype=torch.long)
-    inputs = inputs.to("cuda")
-    attention_mask = attention_mask.to("cuda")    
-    outputs = model.generate(
-        inputs, 
-        attention_mask=attention_mask,
-        max_length=max_length, 
-        top_p=0.9, 
-        temperature=temperature, 
-        do_sample=True, 
-        pad_token_id=tokenizer.eos_token_id
-    )
-    return tokenizer.decode(outputs[0])
-
-def gradio_app():
-    with gr.Blocks() as demo:
-        gr.Markdown(title)
-        prompt = gr.Code(label="Enter your code prompt", value="def print_hello_world():")
-        with gr.Row():
-            temperature = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.5, label="Temperature")
-            max_length = gr.Slider(minimum=100, maximum=1024, step=10, value=450, label="Generate Length")
-        generate_btn = gr.Button("Try✨StarCoder")
-        output = gr.Code(label="✨StarCoder:", lines=40)
-
-        generate_btn.click(
-            fn=generate_text,
-            inputs=[prompt, temperature, max_length],
-            outputs=output
-        )
-
-    demo.launch()
+from huggingface_hub import InferenceClient
+from transformers import pipeline
+from typing import List, Dict, Any
+
+# Initialize the InferenceClient (make sure to add your Hugging Face model)
+client = InferenceClient("your_huggingface_model_name_or_api_key")
+
+def generate_attack(
+    prompt: str,
+    history: List[Dict[str, str]],
+) -> List[str]:
+    """
+    Simulates a Blackhat AI scenario by generating attack strategies and potential impacts.
+    
+    Args:
+        prompt (str): The user's input to the simulator.
+        history (List[Dict]): The user's message history with timestamps.
+    
+    Returns:
+        List[str]: A list of attack responses from the AI.
+    """
+    messages = [
+        {
+            "role": "system",
+            "content": f"Responding to {prompt}..."
+        }
+    ]
+
+    for val in history:
+        if "user" in val:
+            messages.append({"role": "user", "content": val["user"]})
+        if "assistant" in val:
+            messages.append({"role": "assistant", "content": val["assistant"]})
+
+    messages.append({"role": "user", "content": prompt})
+    
+    # Use Hugging Face's Inference API to generate responses
+    generator = pipeline("text-generation", model="your_huggingface_model_name")
+    response = generator(messages[-1]["content"], max_length=100)
+    
+    return [response[0]['generated_text']]  # Return the generated text
+
+demo = gr.Interface(
+    fn=generate_attack,  # or simulate_attack depending on what you want to use
+    inputs=[gr.Textbox()],
+    outputs=gr.Textbox(),
+    title="Blackhat AI Simulator",
+    description="This simulator generates adversarial scenarios, analyzes attack vectors, and provides ethical countermeasures. Use responsibly for cybersecurity training and awareness."
+)
 
 if __name__ == "__main__":
-    gradio_app()
+    demo.launch()