Create app.py

app.py

@@ -0,0 +1,165 @@
+import gradio as gr
+import torch
+
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM
+)
+
+from peft import PeftModel
+
+
+# ==========================================
+# MODEL CONFIG
+# ==========================================
+
+BASE_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
+
+ADAPTER_MODEL = "YOUR_USERNAME/DhanushAI-0.5B"
+
+
+# ==========================================
+# DEVICE
+# ==========================================
+
+if torch.cuda.is_available():
+    DEVICE = "cuda"
+    DTYPE = torch.float16
+else:
+    DEVICE = "cpu"
+    DTYPE = torch.float32
+
+
+print("Device:", DEVICE)
+
+
+# ==========================================
+# TOKENIZER
+# ==========================================
+
+print("Loading tokenizer...")
+
+tokenizer = AutoTokenizer.from_pretrained(
+    BASE_MODEL
+)
+
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+
+# ==========================================
+# BASE MODEL
+# ==========================================
+
+print("Loading base model...")
+
+base_model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL,
+    torch_dtype=DTYPE
+)
+
+base_model.to(DEVICE)
+
+
+# ==========================================
+# YOUR LORA MODEL
+# ==========================================
+
+print("Loading DhanushAI...")
+
+model = PeftModel.from_pretrained(
+    base_model,
+    ADAPTER_MODEL
+)
+
+model.to(DEVICE)
+
+model.eval()
+
+print("DhanushAI loaded!")
+
+
+# ==========================================
+# CHAT FUNCTION
+# ==========================================
+
+def chat(message):
+
+    if not message or not message.strip():
+        return "Please enter a message."
+
+    prompt = f"""You are DhanushAI, a helpful AI assistant.
+
+User: {message}
+
+Assistant:"""
+
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt"
+    )
+
+    inputs = {
+        key: value.to(DEVICE)
+        for key, value in inputs.items()
+    }
+
+    with torch.no_grad():
+
+        outputs = model.generate(
+            **inputs,
+
+            max_new_tokens=200,
+
+            temperature=0.7,
+
+            top_p=0.9,
+
+            do_sample=True,
+
+            repetition_penalty=1.1
+        )
+
+    generated = tokenizer.decode(
+        outputs[0],
+        skip_special_tokens=True
+    )
+
+    if "Assistant:" in generated:
+
+        answer = generated.split(
+            "Assistant:",
+            1
+        )[1].strip()
+
+    else:
+
+        answer = generated.strip()
+
+    return answer
+
+
+# ==========================================
+# GRADIO API
+# ==========================================
+
+demo = gr.Interface(
+
+    fn=chat,
+
+    inputs=gr.Textbox(
+        label="Message",
+        placeholder="Ask DhanushAI something..."
+    ),
+
+    outputs=gr.Textbox(
+        label="DhanushAI"
+    ),
+
+    title="DhanushAI",
+
+    description="My custom AI model"
+)
+
+
+demo.launch()