[Doc] Add the use of the tables in the QuickStart

README.md

@@ -52,9 +52,9 @@ For now, the standalone decoder is open-sourced and fully functional without hav
 
 This model is static, trained on an offline dataset. Future versions may be released to enhance its performance on specialized tasks.
 
-**Quickstart**
+**QuickStart**
 
-Here provides a code snippet with apply_chat_template to show you how to load the tokenizer and model and how to generate contents.
+This code snippet demonstrates how to build a prompt with table information, and shows how to load the tokenizer, load the model, and generate content.
 
 > Note that you need `transformers>=4.37.0` to use `TableGPT2`:
 > ```sh
@@ -64,33 +64,62 @@ Here provides a code snippet with apply_chat_template to show you how to load th
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+# Using pandas to read some structured data
+import pandas as pd
+from io import StringIO
+
+# single table
+EXAMPLE_CSV_CONTENT = """
+"Loss","Date","Score","Opponent","Record","Attendance"
+"Hampton (14–12)","September 25","8–7","Padres","67–84","31,193"
+"Speier (5–3)","September 26","3–1","Padres","67–85","30,711"
+"Elarton (4–9)","September 22","3–1","@ Expos","65–83","9,707"
+"Lundquist (0–1)","September 24","15–11","Padres","67–83","30,774"
+"Hampton (13–11)","September 6","9–5","Dodgers","61–78","31,407"
+"""
+
+csv_file = StringIO(EXAMPLE_CSV_CONTENT)
+df = pd.read_csv(csv_file)
+# Some data preprocessing
+# code
+
 model_name = "tablegpt/TableGPT2-7B"
 
 model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype="auto",
-    device_map="auto"
+    model_name, torch_dtype="auto", device_map="auto"
 )
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-prompt = "Hey, who are you?"
+example_prompt_template = """Given access to several pandas dataframes, write the Python code to answer the user's question.
+
+/*
+"{var_name}.head(5).to_string(index=False)" as follows:
+{df_info}
+*/
+
+Question: {user_question}
+"""
+question = "哪些比赛的战绩达到了40胜40负？"
+
+prompt = example_prompt_template.format(
+    var_name="df",
+    df_info=df.head(5).to_string(index=False),
+    user_question=question,
+)
+
 messages = [
     {"role": "system", "content": "You are a helpful assistant."},
-    {"role": "user", "content": prompt}
+    {"role": "user", "content": prompt},
 ]
 text = tokenizer.apply_chat_template(
-    messages,
-    tokenize=False,
-    add_generation_prompt=True
+    messages, tokenize=False, add_generation_prompt=True
 )
 model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
 
-generated_ids = model.generate(
-    **model_inputs,
-    max_new_tokens=512
-)
+generated_ids = model.generate(**model_inputs, max_new_tokens=512)
 generated_ids = [
-    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+    output_ids[len(input_ids) :]
+    for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
 ]
 
 response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
@@ -121,6 +150,7 @@ For deployment, we recommend using vLLM.
       }'
   
   ```
+  For more details about how to use TableGPT2, please refer to [our repository on GitHub](https://github.com/tablegpt/tablegpt-agent)
 
 **License**