|
|
|
|
|
"""RAG2
|
|
|
|
|
|
Automatically generated by Colab.
|
|
|
|
|
|
Original file is located at
|
|
|
https://colab.research.google.com/drive/1fskQmtugai5co1I64Hv3iAcKQKUKHAzU
|
|
|
"""
|
|
|
|
|
|
from langchain_text_splitters import CharacterTextSplitter
|
|
|
import os
|
|
|
import chromadb
|
|
|
from chromadb.utils import embedding_functions
|
|
|
import sentence_transformers
|
|
|
from sentence_transformers import SentenceTransformer
|
|
|
import ollama
|
|
|
|
|
|
|
|
|
|
|
|
text_splitter2 = CharacterTextSplitter(
|
|
|
separator = " // ",
|
|
|
chunk_size=100,
|
|
|
chunk_overlap=20,
|
|
|
length_function=len,
|
|
|
is_separator_regex=False,
|
|
|
)
|
|
|
|
|
|
final_items = []
|
|
|
|
|
|
directory = r"data/*"
|
|
|
files = os.listdir(directory)
|
|
|
|
|
|
for file in files:
|
|
|
file_path = os.path.join(directory, file)
|
|
|
with open(file_path, 'r') as f:
|
|
|
file_content = f.read()
|
|
|
items = text_splitter2.create_documents([file_content])
|
|
|
final_items.extend(items)
|
|
|
|
|
|
|
|
|
|
|
|
CHROMA_DATA_PATH = r"CHROMA_EMBEDDINGS_PATH"
|
|
|
COLLECTION_NAME = "BioRAG_Collection"
|
|
|
EMBED_MODEL = "all-MiniLM-L6-v2"
|
|
|
client = chromadb.PersistentClient(path = CHROMA_DATA_PATH)
|
|
|
|
|
|
embedding_func = embedding_functions.SentenceTransformerEmbeddingFunction(
|
|
|
model_name=EMBED_MODEL
|
|
|
)
|
|
|
|
|
|
collection = client.create_collection(
|
|
|
name = "BioRAG_Collection",
|
|
|
embedding_function=embedding_func,
|
|
|
metadata={"hnsw:space": "cosine"},
|
|
|
)
|
|
|
|
|
|
documents = []
|
|
|
|
|
|
|
|
|
for item in final_items:
|
|
|
print(item)
|
|
|
prompt = f'Please summarize this segment of Antimony: {item}. The summaries must be clear and concise. For Display Names, provide the value for each variable. Expand mathematical functions into words. Cross reference all parts of the provided context. Explain well without errors and in an easily understandable way. Write in a list format. '
|
|
|
documents5 = ollama.generate(model = "llama3", prompt=prompt)
|
|
|
documents2 = documents5["response"]
|
|
|
documents.append(documents2)
|
|
|
|
|
|
collection.add(
|
|
|
documents = documents,
|
|
|
ids=[f"id{i}" for i in range(len(documents))]
|
|
|
)
|
|
|
|
|
|
|
|
|
while 1==1:
|
|
|
query_text = input("What question would you like to ask BioRAG? If you would like to end the session, please type 'STOP'." )
|
|
|
if query_text == "STOP":
|
|
|
break
|
|
|
query_results = collection.query(
|
|
|
query_texts = query_text,
|
|
|
n_results=5,
|
|
|
)
|
|
|
best_recommendation = query_results['documents']
|
|
|
|
|
|
prompt_template = f"""Use the following pieces of context to answer the question at the end. If you don't know the answer, say so.
|
|
|
|
|
|
This is the piece of context necessary: {best_recommendation}
|
|
|
|
|
|
Cross-reference all pieces of context to define variables and other unknown entities. Calculate mathematical values based on provided matching variables. Remember previous responses if asked a follow up question.
|
|
|
|
|
|
Question: {query_text}
|
|
|
|
|
|
"""
|
|
|
response = ollama.generate(model = "llama3", prompt=prompt_template)
|
|
|
print(response['response']) |
