Update app.py

app.py

@@ -1,39 +1,41 @@
-import gradio as gr
-from huggingface_hub import InferenceClient
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-client = InferenceClient("dexcommunity/dex")  # 👈 Apna Dex model ka path
+# ✅ Load your own model from Hugging Face
+model_id = "dexcommunity/dex"  # 👈 yahan apna model path do
 
-def respond(message, history, system_message, max_tokens, temperature, top_p):
-    prompt = ""
-    for user, bot in history:
-        prompt += f"User: {user}\nDex: {bot}\n"
-    prompt += f"User: {message}\nDex:"
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id)
 
-    response = ""
-    try:
-        for token in client.text_generation(
-            prompt,
-            max_new_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            stream=True,
-        ):
-            response += token
-            yield response
-    except Exception as e:
-        yield f"❌ Error: {str(e)}"
+# Check device (GPU ya CPU)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
-demo = gr.ChatInterface(
-    fn=respond,
-    additional_inputs=[
-        gr.Textbox(value="You are Dex AI 🤖", label="System Message"),
-        gr.Slider(1, 2048, value=512, step=1, label="Max Tokens"),
-        gr.Slider(0.1, 4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p"),
-    ],
-    title="Dex 🛡️ — Your Cybersecurity AI",
-    description="Talk to Dex, your friendly cybersecurity and hacking assistant powered by HuggingFace 🤖"
-)
+# Chat function
+def ask_dex(prompt, max_length=256):
+    input_text = f"User: {prompt}\nDex:"
+    inputs = tokenizer(input_text, return_tensors="pt").to(device)
 
+    output = model.generate(
+        **inputs,
+        max_length=max_length,
+        do_sample=True,
+        top_k=50,
+        top_p=0.9,
+        temperature=0.7,
+        pad_token_id=tokenizer.eos_token_id
+    )
+
+    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
+    reply = decoded.split("Dex:")[-1].strip()
+    return reply
+
+# Example use
 if __name__ == "__main__":
-    demo.launch()
+    while True:
+        user_input = input("You: ")
+        if user_input.lower() in ["exit", "quit"]:
+            break
+        response = ask_dex(user_input)
+        print("Dex:", response)