Update main.py

main.py

@@ -4,62 +4,210 @@ from fastapi.responses import FileResponse, HTMLResponse
 import os
 import io
 
-#import httpcore
-#setattr(httpcore, 'SyncHTTPTransport', 'AsyncHTTPProxy')
+import torch
+from auto_gptq import AutoGPTQForCausalLM
+from langchain import HuggingFacePipeline, PromptTemplate
+from langchain.chains import RetrievalQA
+from langchain.document_loaders import PyPDFDirectoryLoader
+from langchain.embeddings import HuggingFaceInstructEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.vectorstores import Chroma
+from langchain.vectorstores import FAISS
+from transformers import AutoTokenizer, TextStreamer, pipeline
 
-import googletrans
-from googletrans import Translator
-translator = Translator()
-lan = googletrans.LANGUAGES
-keys = list(lan.keys())
-vals = list(lan.values())
+DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-#from gradio_client import Client
 
-#client = Client("physician-ai/speech-to-text")
 
-#print(client.view_api())
-app = FastAPI()
 
 
-@app.post("/translate/")
-async def translate(text,language):
-    return {"translated_text": translator.translate(text,dest=keys[vals.index(language)]).text}
+embeddings = HuggingFaceInstructEmbeddings(
+    model_name="hkunlp/instructor-large", model_kwargs={"device": DEVICE}
+)
+new_db = FAISS.load_local("faiss_index", embeddings)
+
+model_name_or_path = "TheBloke/Llama-2-13B-chat-GPTQ"
+model_basename = "model"
+
+
+tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)
+
+model = AutoGPTQForCausalLM.from_quantized(
+    model_name_or_path,
+    revision="gptq-4bit-128g-actorder_True",
+    model_basename=model_basename,
+    use_safetensors=True,
+    trust_remote_code=True,
+    device=DEVICE,
+    inject_fused_attention=False,
+    quantize_config=None,
+)
+
+#default promts it will work when we don't set the our custom system propts
+DEFAULT_SYSTEM_PROMPT = """
+You are a helpful, respectful and honest assistant. give answer for any questions.
+""".strip()
+
+
+def generate_prompt(prompt: str, system_prompt: str = DEFAULT_SYSTEM_PROMPT) -> str:
+    return f"""
+[INST] <<SYS>>
+{system_prompt}
+<</SYS>>
+{prompt} [/INST]
+""".strip()
+
+# setting the RAG pipeline
+streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+text_pipeline = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=4096,
+    temperature=2,
+    top_p=0.95,
+    repetition_penalty=1.15,
+    streamer=streamer,
+)
+global llm,llm2
+llm = HuggingFacePipeline(pipeline=text_pipeline, model_kwargs={"temperature": 2})
+llm2 = HuggingFacePipeline(pipeline=text_pipeline, model_kwargs={"temperature": 2})
+# when the user query is not related to trained PDF data model will give the response from own knowledge 
+SYSTEM_PROMPT = "give answer from external data's. don't use the provided context"
+
+template = generate_prompt(
+    """
+{context}
+Question: {question}
+""",
+    system_prompt=SYSTEM_PROMPT,
+)
+prompt = PromptTemplate(template=template, input_variables=["context", "question"])
+
+global qa_chain,qa_chain_a
+qa_chain = RetrievalQA.from_chain_type(
+    llm=llm,
+    chain_type="stuff",
+    retriever=new_db.as_retriever(search_kwargs={"k": 2}),
+    return_source_documents=True,
+    chain_type_kwargs={"prompt": prompt},
+)
 
-#@app.post("/speech_to_text/")
-#async def speech_to_text(file: UploadFile = File(...)):
+qa_chain_a = RetrievalQA.from_chain_type(
+    llm=llm2,
+    chain_type="stuff",
+    retriever=new_db.as_retriever(search_kwargs={"k": 2}),
+    return_source_documents=True,
+    chain_type_kwargs={"prompt": prompt},
+)
 
-        # Save the file with a specific name
-        #file_path = "inputvoice.mp3"
+report_prompt_template = """
+this is report format
+Patient Name: [Insert name here]<br>
+Age: [Insert age here]<br>
+sex: [Insert  here]<br>
+Chief Complaint: [insert here]<br>
+History of Present Illness:[insert here]<br>
+Past Medical History: [insert here]<br>
+Medication List: [insert here]<br>
+Social History: [insert here]<br>
+Family History: [insert here]<br>
+Review of Systems: [insert here]<br>
+ICD Code: [insert here]
+convert this bellow details into above format don't add any other details .don't use the provided pdfs data's.\n\n"""
 
-        #with open(file_path, "wb") as f:
-            #f.write(file.file.read())
-        #print("saved")
-        #respond = client.predict(file_path,api_name="/get_stt") 
-        #print(respond.result())
-        #return respond
 
+# 4. prompt sets for ask some defined questions and its will guide the model correct way
+final_question ={
+    8:"Do you have a history of medical conditions, such as allergies, chronic illnesses, or previous surgeries? If so, please provide details.",
+    9:"What medications are you currently taking, including supplements and vitamins?",
+    10:"Can you please Describe Family medical history (particularly close relatives): Does anyone in your immediate family suffer from similar symptoms or health issues?",
+    11:"Can you please Describe Social history: Marital status, occupation, living arrangements, education level, and support system.",
+    12:"Could you describe your symptoms, and have you noticed any changes or discomfort related to your respiratory, cardiovascular, gastrointestinal, or other body systems?"
+}
 
-os.environ["COQUI_TOS_AGREED"] = "1"
-from TTS.api import TTS
+# 1 . basic first prompt for handled the llama in correct like a family physician
+sys = "You are a general family physician.\n\n"
 
-model_names = TTS().list_models()
-m = model_names[0]
-print(model_names)
-global xtts
-xtts = TTS(m, gpu=True)
-#tts.to("cpu") # no GPU or Amd
-xtts.to("cuda")
+# 5 . prommpts for get the diagnosis with ICD code based on the conversation, its will handle unrelated questions also(not related to diagnosis)
+end_sys_prompts = "\n\ngive correct treatment and most related diagnosis with ICD code don't ask any questions. if question is not related to provided data don't give answer from this provided data's"
 
-@app.get("/text-to-speech/")
-def text_to_speech(text,language):
-    global xtts
-    audio_file = 'text_to_speech.wav'
-    if language=="vietnamese":
-        from gtts import gTTS
-        tts = gTTS(text)
-        tts.save(audio_file)
+
+
+def refresh_model():
+        global llm,llm2
+        llm = HuggingFacePipeline(pipeline=text_pipeline, model_kwargs={"temperature": 2})
+        llm2 = HuggingFacePipeline(pipeline=text_pipeline, model_kwargs={"temperature": 2})
+
+        global qa_chain,qa_chain_a
+        qa_chain = RetrievalQA.from_chain_type(
+            llm=llm,
+            chain_type="stuff",
+            retriever=new_db.as_retriever(search_kwargs={"k": 2}),
+            return_source_documents=True,
+            chain_type_kwargs={"prompt": prompt},
+        )
+        
+        qa_chain_a = RetrievalQA.from_chain_type(
+            llm=llm2,
+            chain_type="stuff",
+            retriever=new_db.as_retriever(search_kwargs={"k": 2}),
+            return_source_documents=True,
+            chain_type_kwargs={"prompt": prompt},
+        )
+        print("model refreshed")
+
+app = FastAPI()
+@app.post("/llm_response/")
+async def llm_response(chain,id,mode):
+    id = int(id)
+    global qa_chain,qa_chain_a
+    refresh_model()
+    
+    if id<13:
+        if id>=8:
+            return final_question[id]
+        else:
+            if id<5:
+                # 2 . prompmt control the natural way on question asking based on patient response,symptomps type
+                question = qa_chain(sys+chain +"""\n\nask single small queston to get details based on the patient response,and don't ask 
+                same question again, and don't provide treatment and diagnosis ask next small and short question , 
+                always don't ask same question again and again , always only ask next single small question""")
+
+                    
+            else:
+                # 3. prompt will guide the model to ask yes or no questions based on patient response,symptomps type
+                question = qa_chain(sys+chain +"""\n\nask single small queston to get details based on the patient response,and don't ask 
+                same question again, and don't provide treatment and diagnosis ask next small and short question with yes or no format , 
+                always don't ask same question again and again , always only ask next single small question""")
+            try:
+                if "Patient:" in str(question['result']) or "Patient response:" in str(question['result']):
+                    return str((str(question['result']).split("\n\n")[-1]).split(":")[-1])
+                else:
+                    return str(question['result']).split("\n\n")[1]
+
+            except:
+                if "Patient:" in str(question['result']) or "Patient response:" in str(question['result']):
+                   return str(question['result']).split(":")[-1]
+                else:
+                   return str(question['result'])
+
+    if id==16:
+        diagnosis_and_treatment = qa_chain(sys+chain+end_sys_prompts)
+        diagnosis_and_treatment = str(diagnosis_and_treatment['result'])
+
+        if mode!="h&p":
+            return diagnosis_and_treatment
+        else:
+            report = qa_chain_a(report_prompt_template+sys+chain+"\n\ntreatment & diagnosis with ICD code below\n"+diagnosis_and_treatment)
+            return str(report['result'])
+            
+    result_ex = qa_chain(sys+chain+"""\n\n\nalways give small and single response based on the patient 
+        response. don't give multiline response always give response based on last patient response""")
+    if "Patient:" in str(result_ex['result']) or "Patient response:" in str(result_ex['result']) or "Patient Response" in str(result_ex['result']):
+        return str((str(result_ex['result']).split("\n\n")[-1]).split(":")[-1])
     else:
-        xtts.tts_to_file(text,speaker_wav="input.wav",language=keys[vals.index(language)],file_path=audio_file)
-    return FileResponse(audio_file, media_type='audio/mpeg')
+        return str(result_ex['result']).split("\n\n")[1]
+
+