Spaces:
Runtime error
Runtime error
File size: 2,938 Bytes
2681ef4 c6e8170 2681ef4 2df718a 2681ef4 c6e8170 a3f5e7a 68cfdb8 2681ef4 2df718a 2681ef4 7f5a5e2 |
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 |
# Import necessary libraries
from langchain.document_loaders import DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import SentenceTransformerEmbeddings
from langchain.vectorstores import Chroma
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from langchain_community.llms import HuggingFacePipeline
from langchain.chains.question_answering import load_qa_chain
# Load and process documents
dir = "data"
def load_docs(dir):
loader = DirectoryLoader(dir)
docs = loader.load()
return docs
docs = load_docs(dir)
def split_docs(docs, chunk_size=512, chunk_overlap=20):
text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
split_docs = text_splitter.split_documents(docs)
return split_docs
docs = split_docs(docs)
# Initialize embeddings and vector store
embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
persist_directory = "chroma_db"
vectordb = Chroma.from_documents(docs, embeddings, persist_directory=persist_directory)
vectordb.persist()
new_db = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
def get_similar_docs(query, k=2, score=False):
if score:
similar_docs = new_db.similarity_search_with_score(query, k=k)
else:
similar_docs = new_db.similarity_search(query, k=k)
return similar_docs
# Load LLM model from Hugging Face
# model_name = "HuggingFaceH4/zephyr-7b-beta"
# model = AutoModelForCausalLM.from_pretrained(model_name)
# tokenizer = AutoTokenizer.from_pretrained(model_name)
# model = AutoModelForCausalLM.from_pretrained("gpt2")
# tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama_v1.1")
model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama_v1.1")
text_generation_pipeline = pipeline(
model=model,
tokenizer=tokenizer,
task="text-generation",
temperature=0.2,
do_sample=True,
repetition_penalty=1.1,
return_full_text=True,
max_new_tokens=400,
)
# text_generation_pipeline = pipeline("text-generation", model="bigscience/bloom-1b7")
llm = HuggingFacePipeline(pipeline=text_generation_pipeline)
chain = load_qa_chain(llm, chain_type="stuff")
def get_helpful_answer(text):
# Find the index of "Helpful Answer:"
index = text.find("Helpful Answer:")
# If "Helpful Answer:" is not found, return an empty string
if index == -1:
return ""
# Add the length of "Helpful Answer:" to the index to start from the end of this string
index += len("Helpful Answer:")
# Return the text from this index to the end
return text[index:].strip()
def get_answer(query):
similar_docs = get_similar_docs(query)
answer = chain.run(input_documents=similar_docs, question=query)
answer = get_helpful_answer(answer)
return answer
|