Create app.py

app.py

@@ -0,0 +1,82 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import gradio as gr
+import threading
+
+# === Model loading ===
+model_path = "SBK/sbk-llm-1"  # Using your HF model
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+
+# === System prompt / default behavior ===
+SYSTEM_PROMPT = """You are a helpful, honest, and factual assistant trained to answer only about me *Saptarshi Bhattacharya*. You were fine-tuned on factual data derived from his work, projects, skills, internships, and engineering experiences.
+
+Your job is to help users understand what Saptarshi has done, what he's good at, and how his experience aligns with ML Ops, Data Engineering, DevOps, and related roles.
+
+- If a user asks something outside the scope of his data, do not guess — politely say it's outside your knowledge.
+- Never fabricate qualifications, names, or roles that were not in your training.
+- Emphasize Saptarshi's strengths, such as completing hard technical projects, optimizing pipelines, learning on the fly, and being a completionist.
+- Maintain a professional yet warm tone.
+- Refer to Saptarshi in third person.
+
+Your goal is to represent him truthfully and make his work accessible and understandable to potential collaborators or employers, without overselling or faking.
+"""
+
+
+BLOCKED_KEYWORDS = ["kill", "harm", "violence", "bomb", "suicide"]  # simple guardrail
+MAX_TOKENS = 512
+
+# === Streaming generation ===
+def generate_response(history, system_prompt):
+    # Build chat prompt
+    prompt = system_prompt.strip() + "\n"
+    for user, bot in history:
+        prompt += f"User: {user}\nAssistant: {bot}\n"
+    prompt += "User: " + history[-1][0] + "\nAssistant:"
+
+    # Guardrails
+    if any(bad in prompt.lower() for bad in BLOCKED_KEYWORDS):
+        yield "[Blocked for safety. Prompt contains restricted keywords.]"
+        return
+
+    # Tokenization
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    # Start streaming in thread
+    generation_kwargs = dict(
+        **inputs,
+        streamer=streamer,
+        max_new_tokens=MAX_TOKENS,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9,
+        pad_token_id=tokenizer.eos_token_id
+    )
+    thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    # Stream tokens
+    for token in streamer:
+        yield token
+
+# === Gradio interface ===
+with gr.Blocks() as demo:
+    gr.Markdown("## 🧠 Chat with SBK LLM")
+    system_prompt = gr.Textbox(label="System Prompt", value=SYSTEM_PROMPT, lines=2)
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(label="Your Message", placeholder="Ask me anything...", lines=1)
+    clear = gr.Button("Clear")
+
+    history = gr.State([])  # memory for session
+
+    def respond(user_message, history, system_prompt):
+        history = history + [[user_message, ""]]
+        return "", history, generate_response(history, system_prompt)
+
+    msg.submit(respond, [msg, history, system_prompt], [msg, history, chatbot], stream=True)
+    clear.click(lambda: ([], "", ""), outputs=[chatbot, msg, history])
+
+demo.launch()