Spaces:
Runtime error
Runtime error
| 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 | |
| def mysearch(query: str) -> str: | |
| """Query our own datasets. | |
| """ | |
| rslt = get_text(query) | |
| return '\n'.join(rslt) | |
| def mygreetings(greeting: str) -> str: | |
| """Let us do our greetings | |
| """ | |
| return "how are you?" | |