Update app.py

app.py

@@ -1,117 +1,27 @@
-from pymed import PubMed
-from typing import List
-from haystack import component
-from haystack import Document
-from haystack.components.generators import HuggingFaceTGIGenerator
-from dotenv import load_dotenv
-import os 
-from haystack import Pipeline
-from haystack.components.builders.prompt_builder import PromptBuilder
 import gradio as gr
-import time
-
-# load_dotenv()
-
-# os.environ['HUGGINGFACE_API_KEY'] = os.getenv('HUGGINGFACE_API_KEY')
-
-
-pubmed = PubMed(tool="Haystack2.0Prototype", email="[email protected]")
-
-def documentize(article):
-  return Document(content=article.abstract, meta={'title': article.title, 'keywords': article.keywords})
-
-@component
-class PubMedFetcher():
-
-  @component.output_types(articles=List[Document])
-  def run(self, queries: list[str]):
-    cleaned_queries = queries[0].strip().split('\n')
-
-    articles = []
-    try:
-      for query in cleaned_queries:
-        response = pubmed.query(query, max_results = 1)
-        documents = [documentize(article) for article in response]
-        articles.extend(documents)
-    except Exception as e:
-        print(e)
-        print(f"Couldn't fetch articles for queries: {queries}" )
-    results = {'articles': articles}
-    return results
-
-keyword_llm = HuggingFaceTGIGenerator("liuhaotian/llava-v1.6-mistral-7b")
-keyword_llm.warm_up()
-
-llm = HuggingFaceTGIGenerator("liuhaotian/llava-v1.6-mistral-7b")
-llm.warm_up()
-
-
-keyword_prompt_template = """
-Your task is to convert the following question into 3 keywords that can be used to find relevant medical research papers on PubMed.
-Here is an examples:
-question: "What are the latest treatments for major depressive disorder?"
-keywords:
-Antidepressive Agents
-Depressive Disorder, Major
-Treatment-Resistant depression
----
-question: {{ question }}
-keywords:
-"""
-
-prompt_template = """
-Answer the question truthfully based on the given documents.
-If the documents don't contain an answer, use your existing knowledge base.
-q: {{ question }}
-Articles:
-{% for article in articles %}
-  {{article.content}}
-  keywords: {{article.meta['keywords']}}
-  title: {{article.meta['title']}}
-{% endfor %}
-"""
-
-keyword_prompt_builder = PromptBuilder(template=keyword_prompt_template)
-
-prompt_builder = PromptBuilder(template=prompt_template)
-fetcher = PubMedFetcher()
-
-pipe = Pipeline()
-
-pipe.add_component("keyword_prompt_builder", keyword_prompt_builder)
-pipe.add_component("keyword_llm", keyword_llm)
-pipe.add_component("pubmed_fetcher", fetcher)
-pipe.add_component("prompt_builder", prompt_builder)
-pipe.add_component("llm", llm)
-
-pipe.connect("keyword_prompt_builder.prompt", "keyword_llm.prompt")
-pipe.connect("keyword_llm.replies", "pubmed_fetcher.queries")
-
-pipe.connect("pubmed_fetcher.articles", "prompt_builder.articles")
-pipe.connect("prompt_builder.prompt", "llm.prompt")
-
-def ask(question):
-  output = pipe.run(data={"keyword_prompt_builder":{"question":question},
-                          "prompt_builder":{"question": question},
-                          "llm":{"generation_kwargs": {"max_new_tokens": 500}}})
-  print(question)
-  print(output['llm']['replies'][0])
-  return output['llm']['replies'][0]
-
-# result = ask("How are mRNA vaccines being used for cancer treatment?")
-
-# print(result)
-
-iface = gr.Interface(fn=ask, inputs=gr.Textbox(
-    value="How are mRNA vaccines being used for cancer treatment?"), 
-        outputs="markdown",  
-        title="LLM Augmented Q&A over PubMed Search Engine",
-        description="Ask a question about BioMedical and get an answer from a friendly AI assistant.",
-        examples=[["How are mRNA vaccines being used for cancer treatment?"], 
-                ["Suggest me some Case Studies related to Pneumonia."],
-                ["Tell me about HIV AIDS."],["Suggest some case studies related to Auto Immune Disorders."],
-                ["How to treat a COVID infected Patient?"]],
-    theme=gr.themes.Soft(),
-    allow_flagging="never",)
-
-iface.launch(debug=True)
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Load the quantized model and tokenizer from the Hub
+model = AutoModelForCausalLM.from_pretrained("my-quantized-llava-model")
+tokenizer = AutoTokenizer.from_pretrained("my-quantized-llava-model")
+
+# Define a function to generate a response given an input text and an optional image URL
+def generate_response(text, image_url=None):
+  # Encode the input text and image URL as a single input_ids tensor
+  image_url = "https://upload.wikimedia.org/wikipedia/commons/thumb/3/3a/Roadrunner_Petrochelidon_pyrrhonota.jpg/1200px-Roadrunner_Petrochelidon_pyrrhonota.jpg"
+  if image_url:
+    input_ids = tokenizer(f"{text} <img>{image_url}</img>", return_tensors="pt").input_ids
+  else:
+    input_ids = tokenizer(text, return_tensors="pt").input_ids
+  
+  # Generate a response using beam search with a length penalty of 0.8
+  output_ids = model.generate(input_ids, max_length=256, num_beams=5, length_penalty=0.8)
+
+  # Decode the output_ids tensor into a string
+  output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+  # Return the output text
+  return output_text
+
+# Use the HuggingFaceTGIGenerator class to automatically map inputs and outputs to Gradio components
+gr.Interface(generate_response, gr.HuggingFaceTGIGenerator(model), "text").launch()