aspect2309 commited on
Commit
1fd11df
·
verified ·
1 Parent(s): 79f372a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +109 -0
app.py CHANGED
@@ -0,0 +1,109 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from datasets import Dataset
3
+ from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer, BitsAndBytesConfig
4
+ from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
5
+
6
+ model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
7
+
8
+ bnb_config = BitsAndBytesConfig(
9
+ load_in_4bit=True,
10
+ bnb_4bit_compute_dtype=torch.float16,
11
+ bnb_4bit_use_double_quant=True,
12
+ bnb_4bit_quant_type="nf4",
13
+ )
14
+
15
+ model = AutoModelForCausalLM.from_pretrained(
16
+ model_name,
17
+ quantization_config=bnb_config,
18
+ device_map="auto"
19
+ )
20
+
21
+ tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
22
+ tokenizer.pad_token = tokenizer.eos_token
23
+ model.config.use_cache = False
24
+
25
+ model.gradient_checkpointing_enable()
26
+ model = prepare_model_for_kbit_training(model)
27
+
28
+ lora_config = LoraConfig(
29
+ r=8,
30
+ lora_alpha=32,
31
+ target_modules=["q_proj", "v_proj"],
32
+ lora_dropout=0.05,
33
+ bias="none",
34
+ task_type="CAUSAL_LM"
35
+ )
36
+
37
+ model = get_peft_model(model, lora_config)
38
+
39
+ from datasets import load_dataset
40
+
41
+ # Load the dataset from Hugging Face
42
+ dataset = load_dataset("MakTek/Customer_support_faqs_dataset", split="train")
43
+
44
+ # Print the column names of the dataset
45
+ print("Dataset columns:", dataset.column_names)
46
+
47
+ def format_instruction(example):
48
+ return f"### Instruction:\n{example['question']}\n\n### Response:\n{example['answer']}"
49
+
50
+ dataset = dataset.map(lambda x: {"text": format_instruction(x)})
51
+
52
+ def tokenize_function(example):
53
+ tokenized = tokenizer(example["text"], truncation=True, padding="max_length", max_length=512)
54
+ tokenized["labels"] = tokenized["input_ids"].copy()
55
+ return tokenized
56
+
57
+ tokenized_dataset = dataset.map(tokenize_function, batched=True)
58
+
59
+ training_args = TrainingArguments(
60
+ output_dir="./tinyllama-qlora-support-bot",
61
+ per_device_train_batch_size=2,
62
+ gradient_accumulation_steps=4,
63
+ learning_rate=2e-4,
64
+ logging_dir="./logs",
65
+ num_train_epochs=3,
66
+ logging_steps=10,
67
+ save_total_limit=2,
68
+ save_strategy="epoch",
69
+ bf16=True,
70
+ optim="paged_adamw_8bit"
71
+ )
72
+
73
+ trainer = Trainer(
74
+ model=model,
75
+ args=training_args,
76
+ train_dataset=tokenized_dataset,
77
+ tokenizer=tokenizer
78
+ )
79
+
80
+ trainer.train()
81
+
82
+ model.save_pretrained("tinyllama-qlora-support-bot")
83
+ tokenizer.save_pretrained("tinyllama-qlora-support-bot")
84
+
85
+ from transformers import pipeline
86
+
87
+ pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
88
+
89
+ instruction = "how can i track my order?"
90
+ prompt = f"### Instruction:\n{instruction}\n\n### Response:\n"
91
+
92
+ output = pipe(prompt, max_new_tokens=100)
93
+ print(output[0]['generated_text'])
94
+
95
+
96
+ import gradio as gr
97
+
98
+ def generate_response(instruction):
99
+ prompt = f"### Instruction:\n{instruction}\n\n### Response:\n"
100
+ output = pipe(prompt, max_new_tokens=100, do_sample=True, temperature=0.7)
101
+ return output[0]['generated_text'].replace(prompt, "").strip()
102
+
103
+ gr.Interface(
104
+ fn=generate_response,
105
+ inputs=gr.Textbox(lines=3, placeholder="Ask your customer support question here..."),
106
+ outputs=gr.Textbox(lines=6),
107
+ title="🛠️ Customer Support Chatbot (TinyLlama + QLoRA)",
108
+ description="Ask any support question. Model trained on MakTek/Customer_support_faqs_dataset using TinyLlama 1.1B."
109
+ ).launch()