Hugging Face's logo

Spaces:
Emalawi19
/
Fun-AI-chat
Sleeping

App Files Community
Fun-AI-chat
File size: 2,352 Bytes
f7d0e1b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# =========================
# LOAD MODEL 1
# =========================
model_1_name = "microsoft/DialoGPT-small"
tokenizer_1 = AutoTokenizer.from_pretrained(model_1_name)
model_1 = AutoModelForCausalLM.from_pretrained(model_1_name)

# =========================
# LOAD MODEL 2
# =========================
model_2_name = "microsoft/DialoGPT-medium"
tokenizer_2 = AutoTokenizer.from_pretrained(model_2_name)
model_2 = AutoModelForCausalLM.from_pretrained(model_2_name)

device = "cpu"
model_1.to(device)
model_2.to(device)


# =========================
# GENERATION FUNCTION
# =========================
def generate(model, tokenizer, message, history_ids):
    input_ids = tokenizer.encode(message + tokenizer.eos_token, return_tensors="pt")

    if history_ids is not None:
        bot_input = torch.cat([history_ids, input_ids], dim=-1)
    else:
        bot_input = input_ids

    output = model.generate(
        bot_input,
        max_length=200,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=True,
        top_k=50,
        top_p=0.95,
        temperature=0.9
    )

    reply = tokenizer.decode(
        output[:, bot_input.shape[-1]:][0],
        skip_special_tokens=True
    )

    return reply, output


# =========================
# CHAT LOOP FUNCTION
# =========================
def dual_chat(user_input, history):
    if history is None:
        history = {"msg": "Hello", "h1": None, "h2": None}

    msg = user_input

    chat_log = ""

    for _ in range(3):  # rounds
        chat_log += f"🔵 Model 1: {msg}\n"

        reply2, history["h2"] = generate(model_2, tokenizer_2, msg, history["h2"])
        chat_log += f"🔴 Model 2: {reply2}\n"

        reply1, history["h1"] = generate(model_1, tokenizer_1, reply2, history["h1"])

        msg = reply1

    return chat_log, history


# =========================
# GRADIO UI
# =========================
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 AI vs AI Chat (Hugging Face)")

    input_box = gr.Textbox(label="Start Conversation")
    output_box = gr.Textbox(label="Chat Output")

    state = gr.State()

    btn = gr.Button("Run AI Chat")

    btn.click(
        dual_chat,
        inputs=[input_box, state],
        outputs=[output_box, state]
    )

demo.launch()