Spaces:
Runtime error
Runtime error
File size: 4,351 Bytes
9026f71 | 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 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | from operator import itemgetter
import os
from dotenv import load_dotenv
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from utils import load_docs, split_docs
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_huggingface import ChatHuggingFace
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from langchain_huggingface import HuggingFacePipeline
import warnings
warnings.filterwarnings("ignore", category=UserWarning)
load_dotenv()
HF_TOKEN = os.getenv("HF_TOKEN")
file_path = "data/MedicalBook.pdf"
def downlode_hugging_face_embeddings():
embeddings=HuggingFaceEmbeddings(model_name='sentence-transformers/msmarco-MiniLM-L6-v3')
return embeddings
CHROMA_PATH="chroma_db"
#embeddings = downlode_hugging_face_embeddings()
def create_vectorstore(embeddings):
if os.path.exists(CHROMA_PATH):
print(f"--- Loading existing Vector DB from {CHROMA_PATH}... ---")
# FIX: To LOAD, use Chroma() directly. Do NOT use .from_documents()
return Chroma(
persist_directory=CHROMA_PATH,
embedding_function=embeddings
)
else:
print("Creating a new one...")
documents = load_docs(file_path)
text_chunks = split_docs(documents)
# Use .from_documents ONLY when you have new text_chunks to process
return Chroma.from_documents(
documents=text_chunks,
embedding=embeddings,
persist_directory=CHROMA_PATH
)
# vector_db=create_vectorstore()
# retriever=vector_db.as_retriever(search_kwargs={"k": 3})
def get_llm():
# The correct ID for the 1B Instruct model
model_path = "ibm-granite/granite-3.0-1b-a400m-instruct"
device = -1 # Force CPU
print(f"--- Loading local model: {model_path} ---")
tokenizer = AutoTokenizer.from_pretrained(model_path)
# We use low_cpu_mem_usage to keep the RAM footprint small
model = AutoModelForCausalLM.from_pretrained(
model_path,
torch_dtype=torch.float32,
low_cpu_mem_usage=True,
device_map=None
)
gen_pipeline = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
max_new_tokens=256,
temperature=0.3,
do_sample=True,
return_full_text=False
)
llm = HuggingFacePipeline(pipeline=gen_pipeline)
return ChatHuggingFace(llm=llm)
def get_llm_response(user_input, retriever, llm):
prompt = ChatPromptTemplate.from_messages([
("system", "You are a medical specialist. Use the context to answer. If you don't know, say 'I don't know'."),
("human", "Context: {context}\n\nQuestion: {input}")
])
chain = (
{
"context": itemgetter("input") | retriever,
"input": itemgetter("input"),
}
| prompt
| llm
| StrOutputParser()
)
return chain.invoke({"input": user_input})
def load_system():
"""
Orchestrates the loading of the full RAG pipeline.
This can be called by both FastAPI and Streamlit.
"""
print("--- Initializing Medical AI System ---")
embeddings = downlode_hugging_face_embeddings()
vector_db = create_vectorstore(embeddings)
retriever = vector_db.as_retriever(search_kwargs={"k": 3})
llm = get_llm()
print("--- System Ready ---")
return retriever, llm
if __name__=="__main__":
embeddings = downlode_hugging_face_embeddings()
vector_db=create_vectorstore(embeddings)
retriever=vector_db.as_retriever(search_kwargs={"k": 3})
llm=get_llm()
# print("\n -- LLM Testing -- \n")
# while True:
# User_query=input("User: ")
# if User_query.lower() in ["quit","exit","q"]:
# print("Exiting..")
# break
# response=get_llm_response(User_query, retriever, llm)
# print(f"Assistant: {response}")
|