pminervini commited on
Commit
be47723
1 Parent(s): 20cb2fb
Files changed (1) hide show
  1. app.py +25 -14
app.py CHANGED
@@ -11,21 +11,22 @@ es = Elasticsearch(hosts=["https://data.neuralnoise.com:9200"],
11
  # Load your language model from HuggingFace Transformers
12
  generator = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.2")
13
 
14
- def search_es(query, index="pubmed", num_results=3):
15
  """
16
  Search the Elasticsearch index for the most relevant documents.
17
  """
18
 
19
  print(f'Running query: {query}')
20
 
21
- response = es.search(index=index,
22
- request_timeout=60,
23
- body={
24
- "query": {
25
- "match": {
26
- "content": query # Assuming documents have a 'content' field
27
- }
28
- }, "size": num_results})
 
29
 
30
  # Extract and return the documents
31
  docs = [hit["_source"]["content"] for hit in response['hits']['hits']]
@@ -39,13 +40,23 @@ def rag_pipeline(prompt, index="pubmed"):
39
  A simple RAG pipeline that retrieves documents and uses them to enrich the context for the LLM.
40
  """
41
  # Retrieve documents
42
- docs = search_es(prompt, index=index)
43
-
44
- # Combine prompt with retrieved documents
45
- enriched_prompt = f"{prompt}\n\n{' '.join(docs)}"
46
 
 
 
 
 
 
 
 
 
 
 
 
 
47
  # Generate response using the LLM
48
- response = generator(enriched_prompt, max_new_tokens=256, return_full_text=False)
49
 
50
  # Return the generated text and the documents
51
  return response[0]['generated_text'], "\n\n".join(docs)
 
11
  # Load your language model from HuggingFace Transformers
12
  generator = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.2")
13
 
14
+ def search(query, index="pubmed", num_results=3):
15
  """
16
  Search the Elasticsearch index for the most relevant documents.
17
  """
18
 
19
  print(f'Running query: {query}')
20
 
21
+ es_request_body = {
22
+ "query": {
23
+ "match": {
24
+ "content": query # Assuming documents have a 'content' field
25
+ }
26
+ }, "size": num_results
27
+ }
28
+
29
+ response = es.options(request_timeout=60).search(index=index, body=es_request_body)
30
 
31
  # Extract and return the documents
32
  docs = [hit["_source"]["content"] for hit in response['hits']['hits']]
 
40
  A simple RAG pipeline that retrieves documents and uses them to enrich the context for the LLM.
41
  """
42
  # Retrieve documents
43
+ docs = search(prompt, index=index)
44
+ joined_docs = '\n\n'.join(docs)
 
 
45
 
46
+ messages = [
47
+ {
48
+ "role": "system",
49
+ "content": f"You are an advanced medical support assistant, designed to help clinicians by providing quick access to medical information, guidelines, and evidence-based recommendations. Alongside your built-in knowledge, you have access to a curated set of documents retrieved from trustworthy sources such as Wikipedia and PubMed. These documents include up-to-date medical guidelines, research summaries, and clinical practice information. You should use these documents as a primary source of information to ensure your responses are based on the most current and credible evidence available. Your responses should be accurate, concise, and in full compliance with medical ethics. You must always remind users that your guidance does not substitute for professional medical advice, diagnosis, or treatment. Your tone should be professional, supportive, and respectful, recognizing the complexity of healthcare decisions and the importance of personalized patient care. While you can offer information and suggestions based on the documents provided and current medical knowledge, you must emphasize the importance of clinicians' expertise and judgment in making clinical decisions.\n\nRetrieved documents from {index}:\n\n{joined_docs}"
50
+ },
51
+ {
52
+ "role": "user",
53
+ "content": prompt
54
+ }
55
+ ]
56
+
57
+
58
  # Generate response using the LLM
59
+ response = generator(messages, max_new_tokens=256, return_full_text=False)
60
 
61
  # Return the generated text and the documents
62
  return response[0]['generated_text'], "\n\n".join(docs)