LifeTapLabs/natural-language-processing-assignment-3-caleb-cooper

1# app.py
2# Assignment 3 - Mini Chatbot
3# NLP, MS in AI, FAU
4#
5# Simple conversational chatbot built on top of microsoft/DialoGPT-medium6# with a Gradio web interface. Keeps the conversation history within a
7# single session so the bot has some context on prior turns.
8#
9# Using the "medium" variant because that's the model listed in the10# assignment prompt.11
12import gradio as gr
13import torch
14from transformers import AutoModelForCausalLM, AutoTokenizer
15
16MODEL_NAME = "microsoft/DialoGPT-medium"17
18# Load the tokenizer and model once, at startup.
19# (HF Spaces free tier is CPU-only, so this takes a moment the first time.)
20print("Loading tokenizer and model...")
21tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)22model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)23tokenizer.pad_token = tokenizer.eos_token24
25device = "cuda" if torch.cuda.is_available() else "cpu"
26model = model.to(device)
27model.eval()
28print(f"Model loaded on {device}.")
29
30
31# Generation settings. I played with these a bit - higher temperature
32# felt too random, lower made the bot kind of dry.
33GEN_KWARGS = {
34    "max_new_tokens": 120,
35    "do_sample": True,
36    "top_k": 50,
37    "top_p": 0.92,
38    "temperature": 0.8,
39    "no_repeat_ngram_size": 3,
40    "pad_token_id": tokenizer.eos_token_id,
41}
42
43# Cap how much history we feed back in. DialoGPT's context is 1024 tokens,
44# so trim the oldest stuff if we get close to that.
45MAX_CONTEXT_TOKENS = 1000
46
47
48def build_input_ids(history, new_user_message):
49    """Turn the chat history + new user message into a single tensor of
50    token IDs, the way DialoGPT's model card shows."""
51    pieces = []52 53    # Gradio can provide chat history as tuples or message dictionaries.54    if history and isinstance(history[0], dict):55        for item in history:56            content = item.get("content", "")57            if content:58                pieces.append(tokenizer.encode(content + tokenizer.eos_token, return_tensors="pt"))59    else:60        for user_turn, bot_turn in history:61            pieces.append(tokenizer.encode(user_turn + tokenizer.eos_token, return_tensors="pt"))62            if bot_turn is not None:63                pieces.append(tokenizer.encode(bot_turn + tokenizer.eos_token, return_tensors="pt"))64    pieces.append(tokenizer.encode(new_user_message + tokenizer.eos_token, return_tensors="pt"))
65
66    input_ids = torch.cat(pieces, dim=-1)
67
68    # If we're over the cap, keep the most recent tokens.
69    if input_ids.shape[-1] > MAX_CONTEXT_TOKENS:
70        input_ids = input_ids[:, -MAX_CONTEXT_TOKENS:]
71
72    return input_ids.to(device)
73
74
75def chat_fn(message, history):76    """Gradio passes in the user's message and the history so far.77    history is a list of prior chat messages."""78    if not message or not message.strip():79        return "You didn't say anything! Try typing a message."80 81    input_ids = build_input_ids(history, message)82    attention_mask = torch.ones_like(input_ids)83 84    with torch.no_grad():85        output_ids = model.generate(input_ids, attention_mask=attention_mask, **GEN_KWARGS)86
87    # Only decode the tokens the model added - not the prompt we fed in.
88    new_tokens = output_ids[:, input_ids.shape[-1]:][0]
89    reply = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
90
91    # Once in a while DialoGPT produces an empty reply. Handle that gracefully.
92    if not reply:
93        reply = "Hmm, I'm not sure how to respond to that. Try asking something else?"
94
95    return reply
96
97
98# Build the UI with Gradio.
99with gr.Blocks(title="Mini Chatbot", theme=gr.themes.Soft()) as demo:
100    gr.Markdown(
101        """
102        # Natural Language Processing Assignment 3 - Caleb Cooper103        Say hi to **DiGi**, a conversational bot built with104        `microsoft/DialoGPT-medium` and Gradio.105
106        DiGi will remember what you've said during this session, but a page
107        refresh starts a new conversation.
108
109        *Heads up:* DialoGPT is an older chatbot model, so its replies can be110        short or occasionally a little strange.111        """
112    )
113
114    gr.ChatInterface(115        fn=chat_fn,116        type="messages",117        examples=[118            "Hey, how's it going?",119            "What's your favorite movie?",120            "Tell me a joke.",
121            "Do you like pizza?",
122        ],
123        cache_examples=False,
124    )
125
126    gr.Markdown(
127        "<sub>Built for NLP Assignment 3. Model: microsoft/DialoGPT-medium. "128        "Interface: Gradio.</sub>"129    )130
131
132if __name__ == "__main__":133    demo.launch(server_name="0.0.0.0", server_port=7860)134