Spaces:
Runtime error
Runtime error
File size: 2,186 Bytes
272ec4b |
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 |
from langchain.agents import tool
from torch import tensor as torch_tensor
from datasets import load_dataset
from sentence_transformers import SentenceTransformer, CrossEncoder, util
"""# import models"""
bi_encoder = SentenceTransformer(
'sentence-transformers/multi-qa-MiniLM-L6-cos-v1')
bi_encoder.max_seq_length = 256 # Truncate long passages to 256 tokens
# The bi-encoder will retrieve top_k documents. We use a cross-encoder, to re-rank the results list to improve the quality
cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
"""# import datasets"""
dataset = load_dataset("gfhayworth/wiki_mini", split='train')
mypassages = list(dataset.to_pandas()['psg'])
dataset_embed = load_dataset("gfhayworth/wiki_mini_embed", split='train')
dataset_embed_pd = dataset_embed.to_pandas()
mycorpus_embeddings = torch_tensor(dataset_embed_pd.values)
def search(query, top_k=20, top_n=1):
question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
hits = util.semantic_search(
question_embedding, mycorpus_embeddings, top_k=top_k)
hits = hits[0] # Get the hits for the first query
##### Re-Ranking #####
cross_inp = [[query, mypassages[hit['corpus_id']]] for hit in hits]
cross_scores = cross_encoder.predict(cross_inp)
# Sort results by the cross-encoder scores
for idx in range(len(cross_scores)):
hits[idx]['cross-score'] = cross_scores[idx]
hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
predictions = hits[:top_n]
return predictions
# for hit in hits[0:3]:
# print("\t{:.3f}\t{}".format(hit['cross-score'], mypassages[hit['corpus_id']].replace("\n", " ")))
def get_text(qry):
# predictions = greg_search(qry)
predictions = search(qry)
prediction_text = []
for hit in predictions:
prediction_text.append("{}".format(mypassages[hit['corpus_id']]))
return prediction_text
@tool
def mysearch(query: str) -> str:
"""Query our own datasets.
"""
rslt = get_text(query)
return '\n'.join(rslt)
@tool
def mygreetings(greeting: str) -> str:
"""Let us do our greetings
"""
return "how are you?"
|