txtai：基于 Transformer 的人工智能搜索引擎

pip install txtai

import numpy as npfrom txtai.embeddings import Embeddings# Create embeddings model, backed by sentence-transformers & transformersembeddings = Embeddings({"method": "transformers", "path": "sentence-transformers/bert-base-nli-mean-tokens"})sections = ["US tops 5 million confirmed virus cases","Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg","Beijing mobilises invasion craft along coast as Taiwan tensions escalate","The National Park Service warns against sacrificing slower friends in a bear attack","Maine man wins $1M from $25 lottery ticket","Make huge profits without work, earn up to $100,000 a day"]print("%-20s %s" % ("Query", "Best Match"))print("-" * 50)for query in ("feel good story", "climate change", "health", "war", "wildlife", "asia","north america", "dishonest junk"):# Get index of best section that best matches queryuid = np.argmax(embeddings.similarity(query, sections))print("%-20s %s" % (query, sections[uid]))

# Create an index for the list of sectionsembeddings.index([(uid, text, None) for uid, text in enumerate(sections)])print("%-20s %s" % ("Query", "Best Match"))print("-" * 50)# Run an embeddings search for each queryfor query in ("feel good story", "climate change", "health", "war", "wildlife", "asia","north america", "dishonest junk"):# Extract uid of first result# search result format: (uid, score)uid = embeddings.search(query, 1)[0][0]# Print sectionprint("%-20s %s" % (query, sections[uid]))

embeddings.save("index")embeddings = Embeddings()embeddings.load("index")uid = embeddings.search("climate change", 1)[0][0]print(sections[uid])

from txtai.embeddings import Embeddingsfrom txtai.extractor import Extractor# Create embeddings model, backed by sentence-transformers & transformersembeddings = Embeddings({"method": "transformers", "path": "sentence-transformers/bert-base-nli-mean-tokens"})# Create extractor instanceextractor = Extractor(embeddings, "distilbert-base-cased-distilled-squad")

sections = ["Giants hit 3 HRs to down Dodgers","Giants 5 Dodgers 4 final","Dodgers drop Game 2 against the Giants, 5-4","Blue Jays 2 Red Sox 1 final","Red Sox lost to the Blue Jays, 2-1","Blue Jays at Red Sox is over. Score: 2-1","Phillies win over the Braves, 5-0","Phillies 5 Braves 0 final","Final: Braves lose to the Phillies in the series opener, 5-0","Final score: Flyers 4 Lightning 1","Flyers 4 Lightning 1 final","Flyers win 4-1"]# Add unique id to each section to assist with qa extractionsections = [(uid, section) for uid, section in enumerate(sections)]questions = ["What team won the game?", "What was score?"]execute = lambda query: extractor(sections, [(question, query, question, False) for question in questions])for query in ["Red Sox - Blue Jays", "Phillies - Braves", "Dodgers - Giants", "Flyers - Lightning"]:print("----", query, "----")for answer in execute(query):print(answer)print()# Ad-hoc questionsquestion = "What hockey team won?"print("----", question, "----")print(extractor(sections, [(question, question, question, False)]))