feat(tableQA): add huggingface table-QA
Browse files
multi-agents-analysis/table-QA-huggingface.py
ADDED
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from transformers import pipeline
|
2 |
+
import pandas as pd
|
3 |
+
|
4 |
+
# prepare table + question
|
5 |
+
data = {"Actors": ["Brad Pitt", "Leonardo Di Caprio",
|
6 |
+
"George Clooney"], "Number of movies": ["87", "53", "69"]}
|
7 |
+
# table = pd.DataFrame.from_dict(data)
|
8 |
+
query = "how many movies does Leonardo Di Caprio have?"
|
9 |
+
# data = {
|
10 |
+
# "year": [1896, 1900, 1904, 2004, 2008, 2012],
|
11 |
+
# "city": ["athens", "paris", "st. louis", "athens", "beijing", "london"]
|
12 |
+
# }
|
13 |
+
tablle = pd.DataFrame.from_dict(data)
|
14 |
+
# query = "when london had the olympics game?"
|
15 |
+
|
16 |
+
|
17 |
+
# pipeline model
|
18 |
+
# Note: you must to install torch-scatter first.
|
19 |
+
tqa = pipeline(task="table-question-answering",
|
20 |
+
model="google/tapas-large-finetuned-wtq")
|
21 |
+
|
22 |
+
# result
|
23 |
+
result = tqa(table=tablle, query=query)
|
24 |
+
print(f"Result: {result}")
|
25 |
+
|
26 |
+
print(f"\nResponse: {result['cells'][0]}")
|
27 |
+
# 53
|