Spaces:
Paused
Paused
File size: 2,938 Bytes
6aabcd0 |
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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
# import csv
# # Load sample data (a restaurant menu of items)
# with open('./data.csv') as file:
# lines = csv.reader(file)
# # Store the name of the menu items in this array. In Chroma, a "document" is a string i.e. name, sentence, paragraph, etc.
# documents = []
# # Store the corresponding menu item IDs in this array.
# metadatas = []
# # Each "document" needs a unique ID. This is like the primary key of a relational database. We'll start at 1 and increment from there.
# ids = []
# id = 1
# # Loop thru each line and populate the 3 arrays.
# for i, line in enumerate(lines):
# if i==0:
# # Skip the first row (the column headers)
# continue
# documents.append(line[0])
# metadatas.append({"item_id": line[1]})
# ids.append(str(id))
# id+=1
import chromadb
from chromadb.utils import embedding_functions
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from transformers import pipeline
from langchain.llms import HuggingFacePipeline
import torch
# Instantiate chromadb instance. Data is stored in memory only.
# chroma_client = chromadb.Client()
# Instantiate chromadb instance. Data is stored on disk (a folder named 'my_vectordb' will be created in the same folder as this file).
chroma_client = chromadb.PersistentClient(path="vector_db")
# Select the embedding model to use.
# List of model names can be found here https://www.sbert.net/docs/pretrained_models.html
sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-mpnet-base-v2")
# Use this to delete the database
# chroma_client.delete_collection(name="my_collection")
# Create the collection, aka vector database. Or, if database already exist, then use it. Specify the model that we want to use to do the embedding.
collection = chroma_client.get_or_create_collection(name="my_collection", embedding_function=sentence_transformer_ef)
# collection.add(
# documents=documents,
# metadatas=metadatas,
# ids=ids
# )
results = collection.query(
query_texts=["director"],
n_results=1,
include=['documents', 'distances', 'metadatas']
)
print(results['metadatas'])
tokenizer = AutoTokenizer.from_pretrained("MBZUAI/LaMini-T5-738M")
model = AutoModelForSeq2SeqLM.from_pretrained("MBZUAI/LaMini-T5-738M")
pipe = pipeline(
"text2text-generation",
model=model,
tokenizer=tokenizer,
max_length=512
)
local_llm = HuggingFacePipeline(pipeline=pipe)
context = results['documents'][0][0]
question = "director job"
l = f"""
Use the following pieces of context to answer the question at the end . If you don't know the answer, just say that you don't know, don't try to make up an answer.
{context}
Question: {question}
Helpful Answer:
"""
print(local_llm(l)) |