update

app.py

@@ -1,6 +1,8 @@
 import os
 import gradio as gr
 
+import vllm
+
 import torch
 from transformers import pipeline, StoppingCriteria, StoppingCriteriaList, MaxTimeCriteria, AutoTokenizer, AutoModelForCausalLM, PreTrainedTokenizer
 from openai import OpenAI
@@ -28,11 +30,6 @@ class MultiTokenEOSCriteria(StoppingCriteria):
         return False not in self.done_tracker
 
 
-# Connect to Elasticsearch
-es = Elasticsearch(hosts=["https://data.neuralnoise.com:9200"],
-                   basic_auth=('elastic', os.environ['ES_PASSWORD']),
-                   verify_certs=False, ssl_show_warn=False)
-
 def search(query, index="pubmed", num_docs=3):
     """
     Search the Elasticsearch index for the most relevant documents.
@@ -48,6 +45,12 @@ def search(query, index="pubmed", num_docs=3):
                 }
             }, "size": num_docs
         }
+
+        # Connect to Elasticsearch
+        es = Elasticsearch(hosts=["https://data.neuralnoise.com:9200"],
+                           basic_auth=('elastic', os.environ['ES_PASSWORD']),
+                           verify_certs=False, ssl_show_warn=False)
+
         response = es.options(request_timeout=60).search(index=index, body=es_request_body)
         # Extract and return the documents
         docs = [hit["_source"]["content"] for hit in response['hits']['hits']]
@@ -55,6 +58,31 @@ def search(query, index="pubmed", num_docs=3):
 
     return docs
 
+def analyse(text: str) -> str:
+    model = vllm.LLM(model="fava-uw/fava-model") 
+    sampling_params = vllm.SamplingParams(temperature=0, top_p=1.0, max_tokens=500)
+    outputs = model.generate(text, sampling_params)
+    outputs = [it.outputs[0].text for it in outputs]
+    output = outputs[0].replace("<mark>", "<span style='color: green; font-weight: bold;'> ")
+    output = output.replace("</mark>", " </span>")
+    output = output.replace("<delete>", "<span style='color: red; text-decoration: line-through;'>")
+    output = output.replace("</delete>", "</span>")
+    output = output.replace("<entity>", "<span style='background-color: #E9A2D9; border-bottom: 1px dotted;'>entity</span>")
+    output = output.replace("<relation>", "<span style='background-color: #F3B78B; border-bottom: 1px dotted;'>relation</span>")
+    output = output.replace("<contradictory>", "<span style='background-color: #FFFF9B; border-bottom: 1px dotted;'>contradictory</span>")
+    output = output.replace("<unverifiable>", "<span style='background-color: #D3D3D3; border-bottom: 1px dotted;'>unverifiable</span><u>")
+    output = output.replace("<invented>", "<span style='background-color: #BFE9B9; border-bottom: 1px dotted;'>invented</span>")
+    output = output.replace("<subjective>", "<span style='background-color: #D3D3D3; border-bottom: 1px dotted;'>subjective</span><u>")
+    output = output.replace("</entity>", "")
+    output = output.replace("</relation>", "")
+    output = output.replace("</contradictory>", "")
+    output = output.replace("</unverifiable>", "</u>")
+    output = output.replace("</invented>", "")
+    output = output.replace("</subjective>", "</u>")
+    output = output.replace("Edited:", "")
+    return f'<div style="font-weight: normal;">{output}</div>'
+
+
 def rag_pipeline(prompt, index="pubmed", num_docs=3, model_name="HuggingFaceH4/zephyr-7b-beta"):
     """
     A simple RAG pipeline that retrieves documents and uses them to enrich the context for the LLM.
@@ -118,8 +146,10 @@ def rag_pipeline(prompt, index="pubmed", num_docs=3, model_name="HuggingFaceH4/z
         print('HF_RESPONSE', hf_response)
         response = hf_response[0]['generated_text']
     
+    analysed_response = analyse(response)
+
     # Return the generated text and the documents
-    return response, joined_docs
+    return response, analysed_response, joined_docs
 
 # Create the Gradio interface
 iface = gr.Interface(fn=rag_pipeline,
@@ -130,7 +160,8 @@ iface = gr.Interface(fn=rag_pipeline,
                          gr.Dropdown(label="Model", choices=["HuggingFaceH4/zephyr-7b-beta", "meta-llama/Llama-2-7b-chat-hf", "meta-llama/Llama-2-13b-chat-hf", "meta-llama/Llama-2-70b-chat-hf", "openai/gpt-3.5-turbo"], value="HuggingFaceH4/zephyr-7b-beta")
                      ],
                      outputs=[
-                         gr.Textbox(label="Generated Text"),
+                         gr.Textbox(label="Generated Answer"),
+                         gr.Textbox(label="Analysed Answer"),
                          gr.Textbox(label="Retrieved Documents")
                      ],
                      description="Retrieval-Augmented Generation Pipeline")