Update app.py

app.py

@@ -1,13 +1,12 @@
 import torch
 import gradio as gr
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 from pydub import AudioSegment
 from sentence_transformers import SentenceTransformer, util
 import spacy
 spacy.cli.download("en_core_web_sm")
 import json
 from faster_whisper import WhisperModel
-import ollama
 
 # Audio conversion from MP4 to MP3
 def convert_mp4_to_mp3(mp4_path, mp3_path):
@@ -34,8 +33,8 @@ def load_faster_whisper():
 nlp = spacy.load("en_core_web_sm")
 embedder = SentenceTransformer("all-MiniLM-L6-v2")
 
-tokenizer = AutoTokenizer.from_pretrained("Mahalingam/DistilBart-Med-Summary")
-model = AutoModelForSeq2SeqLM.from_pretrained("Mahalingam/DistilBart-Med-Summary")
+tokenizer = AutoTokenizer.from_pretrained("aws-prototyping/MegaBeam-Mistral-7B-512k")
+model = AutoModelForCausalLM.from_pretrained("aws-prototyping/MegaBeam-Mistral-7B-512k")
 
 summarizer = pipeline("summarization", model=model, tokenizer=tokenizer)
 
@@ -47,12 +46,14 @@ soap_prompts = {
 }
 soap_embeddings = {section: embedder.encode(prompt, convert_to_tensor=True) for section, prompt in soap_prompts.items()}
 
-# Ollama Llama 2 Model Query function
-def ollama_query(user_prompt, soap_note):
+# Query function for MegaBeam-Mistral-7B
+def megabeam_query(user_prompt, soap_note):
     combined_prompt = f"User Instructions:\n{user_prompt}\n\nContext:\n{soap_note}"
     try:
-        response = ollama.chat(model="llama2:7b-uncensored", messages=[{"role": "user", "content": combined_prompt}])
-        return response['text']
+        inputs = tokenizer(combined_prompt, return_tensors="pt")
+        outputs = model.generate(**inputs, max_length=512)
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response
     except Exception as e:
         return f"Error generating response: {e}"
 
@@ -139,7 +140,7 @@ def process_file(mp4_file, user_prompt):
     soap_note = soap_analysis(transcription)
     print("SOAP Notes: ", soap_note)
 
-    template_output = ollama_query(user_prompt, soap_note)
+    template_output = megabeam_query(user_prompt, soap_note)
     print("Template: ", template_output)
 
     json_output = convert_to_json(template_output)
@@ -151,7 +152,7 @@ def process_text(text, user_prompt):
     soap_note = soap_analysis(text)
     print(soap_note)
 
-    template_output = ollama_query(user_prompt, soap_note)
+    template_output = megabeam_query(user_prompt, soap_note)
     print(template_output)
     json_output = convert_to_json(template_output)
 
@@ -170,7 +171,7 @@ def launch_gradio():
                 ],
                 outputs=[
                     gr.Textbox(label="SOAP Note"),
-                    gr.Textbox(label="Generated Template from Llama 2"),
+                    gr.Textbox(label="Generated Template from MegaBeam-Mistral-7B"),
                     gr.Textbox(label="JSON Output"),
                 ],
             )
@@ -183,7 +184,7 @@ def launch_gradio():
                 ],
                 outputs=[
                     gr.Textbox(label="SOAP Note"),
-                    gr.Textbox(label="Generated Template from Llama 2"),
+                    gr.Textbox(label="Generated Template from MegaBeam-Mistral-7B"),
                     gr.Textbox(label="JSON Output"),
                 ],
             )