File size: 2,543 Bytes
5c8c783
3306e69
5c8c783
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3306e69
5c8c783
 
3306e69
 
5c8c783
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3306e69
5c8c783
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3306e69
5c8c783
 
 
 
 
3306e69
 
 
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
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel

# -------------------------------------------------
# Configuration
# -------------------------------------------------
BASE_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
LORA_REPO = "Srikanthgoud7/qwen2.5-customer-support-lora"  

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# -------------------------------------------------
# Load tokenizer & model
# -------------------------------------------------
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    device_map="auto",
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
)

model = PeftModel.from_pretrained(base_model, LORA_REPO)
model.eval()

# -------------------------------------------------
# Chat function (DYNAMIC INPUT)
# -------------------------------------------------
def chat(user_input, history):
    """
    user_input: current user message
    history: list of (user, assistant) tuples
    """

    # Build conversation prompt
    conversation = ""
    for u, a in history:
        conversation += f"User: {u}\nAssistant: {a}\n"
    conversation += f"User: {user_input}\nAssistant:"

    inputs = tokenizer(conversation, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=200,
            temperature=0.7,
            do_sample=True
        )

    response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Extract only the assistant's latest reply
    response = response.split("Assistant:")[-1].strip()

    history.append((user_input, response))
    return history, history

# -------------------------------------------------
# Gradio Chat UI
# -------------------------------------------------
with gr.Blocks() as demo:
    gr.Markdown("## 🤖 Customer Support Chatbot (LoRA + Qwen2.5)")

    chatbot = gr.Chatbot()
    user_input = gr.Textbox(
        placeholder="Type your query here (e.g., Can you cancel my order?)",
        label="Your Message"
    )
    state = gr.State([])

    send_btn = gr.Button("Send")
    clear_btn = gr.Button("Clear Chat")

    send_btn.click(
        fn=chat,
        inputs=[user_input, state],
        outputs=[chatbot, state]
    )

    clear_btn.click(
        fn=lambda: ([], []),
        inputs=None,
        outputs=[chatbot, state]
    )

if __name__ == "__main__":
    demo.launch()