JerameeUC
1st
732e77c
raw
history blame
1.4 kB
import os
import gradio as gr
from transformers import pipeline
MODEL_NAME = os.getenv("HF_MODEL_GENERATION", "distilgpt2")
_pipe = None
def _get_pipe():
global _pipe
if _pipe is None:
_pipe = pipeline("text-generation", model=MODEL_NAME)
return _pipe
def chat_fn(message, max_new_tokens=128, temperature=0.8, top_p=0.95):
message = (message or "").strip()
if not message:
return "Please type something!"
pipe = _get_pipe()
out = pipe(
message,
max_new_tokens=int(max_new_tokens),
do_sample=True,
temperature=float(temperature),
top_p=float(top_p),
pad_token_id=50256
)
return out[0]["generated_text"]
with gr.Blocks(title="Agentic-Chat-bot") as demo:
gr.Markdown("# 🤖 Agentic Chat Bot\nGradio + Transformers demo")
prompt = gr.Textbox(label="Prompt", placeholder="Ask me anything…", lines=4)
out = gr.Textbox(label="Response", lines=6)
max_new = gr.Slider(32, 512, 128, 1, label="Max new tokens")
temp = gr.Slider(0.1, 1.5, 0.8, 0.05, label="Temperature")
topp = gr.Slider(0.1, 1.0, 0.95, 0.05, label="Top-p")
btn = gr.Button("Send")
btn.click(chat_fn, [prompt, max_new, temp, topp], out)
prompt.submit(chat_fn, [prompt, max_new, temp, topp], out)
if __name__ == "__main__":
demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))