Update app.py

app.py

@@ -1,16 +1,12 @@
 import torch
 import gradio as gr
-from transformers import AutoTokenizer
-from gptq import GPTQForCausalLM  # GPTQ model handler
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from pydub import AudioSegment
 from sentence_transformers import SentenceTransformer, util
 import spacy
-spacy.cli.download("en_core_web_sm")
-from subprocess import Popen, PIPE
 import json
 from faster_whisper import WhisperModel
 
-
 # Audio conversion from MP4 to MP3
 def convert_mp4_to_mp3(mp4_path, mp3_path):
     try:
@@ -19,7 +15,6 @@ def convert_mp4_to_mp3(mp4_path, mp3_path):
     except Exception as e:
         raise RuntimeError(f"Error converting MP4 to MP3: {e}")
 
-
 # Check if CUDA is available for GPU acceleration
 if torch.cuda.is_available():
     device = "cuda"
@@ -28,36 +23,21 @@ else:
     device = "cpu"
     compute_type = "int8"
 
-
 # Load Faster Whisper Model for transcription
 def load_faster_whisper():
     model = WhisperModel("deepdml/faster-whisper-large-v3-turbo-ct2", device=device, compute_type=compute_type)
     return model
 
-
-# Load GPTQ Mistral-7B model
-def load_mistral_model():
-    model_name = "TheBloke/Mistral-7B-Instruct-v0.2-GPTQ"
-    
-    # Load the tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    
-    # Load the GPTQ model
-    model = GPTQForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
-
-    return model, tokenizer
-
-
 # Load NLP model and other helpers
 nlp = spacy.load("en_core_web_sm")
 embedder = SentenceTransformer("all-MiniLM-L6-v2")
 
+# Load Summarizer Model (DistilBart-Med-Summary)
 tokenizer = AutoTokenizer.from_pretrained("Mahalingam/DistilBart-Med-Summary")
 model = AutoModelForSeq2SeqLM.from_pretrained("Mahalingam/DistilBart-Med-Summary")
 
 summarizer = pipeline("summarization", model=model, tokenizer=tokenizer)
 
-
 soap_prompts = {
     "subjective": "Personal reports, symptoms described by patients, or personal health concerns. Details reflecting individual symptoms or health descriptions.",
     "objective": "Observable facts, clinical findings, professional observations, specific medical specialties, and diagnoses.",
@@ -66,19 +46,31 @@ soap_prompts = {
 }
 soap_embeddings = {section: embedder.encode(prompt, convert_to_tensor=True) for section, prompt in soap_prompts.items()}
 
+# Load LLaMA 7B model and tokenizer
+def load_llama_model():
+    tokenizer = AutoTokenizer.from_pretrained("huggyllama/llama-7b")
+    model = AutoModelForCausalLM.from_pretrained("huggyllama/llama-7b", torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
+    model.to(device)
+    return model, tokenizer
+
+# Initialize LLaMA model
+llama_model, llama_tokenizer = load_llama_model()
 
-# Llama query function (same as before)
-def llama_query(user_prompt, soap_note, model="llama3.2"):
+# Query function for LLaMA
+def llama_query(user_prompt, soap_note):
     combined_prompt = f"User Instructions:\n{user_prompt}\n\nContext:\n{soap_note}"
     try:
-        process = Popen(['ollama', 'run', model], stdin=PIPE, stdout=PIPE, stderr=PIPE, text=True, encoding='utf-8')
-        stdout, stderr = process.communicate(input=combined_prompt)
-        if process.returncode != 0:
-            return f"Error: {stderr.strip()}"
-        return stdout.strip()
+        inputs = llama_tokenizer(combined_prompt, return_tensors="pt", truncation=True, max_length=4096).to(device)
+        outputs = llama_model.generate(
+            inputs["input_ids"],
+            max_length=512,
+            temperature=0.7,
+            num_beams=4,
+            no_repeat_ngram_size=3
+        )
+        return llama_tokenizer.decode(outputs[0], skip_special_tokens=True)
     except Exception as e:
-        return f"Unexpected error: {str(e)}"
-
+        return f"Error generating response: {e}"
 
 # Convert the response to JSON format
 def convert_to_json(template):
@@ -96,7 +88,6 @@ def convert_to_json(template):
     except Exception as e:
         return f"Error converting to JSON: {e}"
 
-
 # Transcription using Faster Whisper
 def transcribe_audio(mp4_path):
     try:
@@ -112,13 +103,11 @@ def transcribe_audio(mp4_path):
     except Exception as e:
         return f"Error during transcription: {e}"
 
-
 # Classify the sentence to the correct SOAP section
 def classify_sentence(sentence):
     similarities = {section: util.pytorch_cos_sim(embedder.encode(sentence), soap_embeddings[section]) for section in soap_prompts.keys()}
     return max(similarities, key=similarities.get)
 
-
 # Summarize the section if it's too long
 def summarize_section(section_text):
     if len(section_text.split()) < 50:
@@ -134,7 +123,6 @@ def summarize_section(section_text):
     )
     return tokenizer.decode(summary_ids[0], skip_special_tokens=True)
 
-
 # Analyze the SOAP content and divide into sections
 def soap_analysis(text):
     doc = nlp(text)
@@ -150,7 +138,6 @@ def soap_analysis(text):
 
     return format_soap_output(soap_note)
 
-
 # Format the SOAP note output
 def format_soap_output(soap_note):
     return (
@@ -160,7 +147,6 @@ def format_soap_output(soap_note):
         f"Plan:\n{soap_note['plan']}\n"
     )
 
-
 # Process file function for audio to SOAP
 def process_file(mp4_file, user_prompt):
     transcription = transcribe_audio(mp4_file.name)
@@ -176,7 +162,6 @@ def process_file(mp4_file, user_prompt):
 
     return soap_note, template_output, json_output
 
-
 # Process text function for text input to SOAP
 def process_text(text, user_prompt):
     soap_note = soap_analysis(text)
@@ -188,7 +173,6 @@ def process_text(text, user_prompt):
 
     return soap_note, template_output, json_output
 
-
 # Launch the Gradio interface
 def launch_gradio():
     with gr.Blocks(theme=gr.themes.Default()) as demo:
@@ -202,7 +186,7 @@ def launch_gradio():
                 ],
                 outputs=[
                     gr.Textbox(label="SOAP Note"),
-                    gr.Textbox(label="Generated Template from Mistral-7B Instruct"),
+                    gr.Textbox(label="Generated Template from LLaMA"),
                     gr.Textbox(label="JSON Output"),
                 ],
             )
@@ -215,13 +199,12 @@ def launch_gradio():
                 ],
                 outputs=[
                     gr.Textbox(label="SOAP Note"),
-                    gr.Textbox(label="Generated Template from Mistral-7B Instruct"),
+                    gr.Textbox(label="Generated Template from LLaMA"),
                     gr.Textbox(label="JSON Output"),
                 ],
             )
     demo.launch(share=True, debug=True)
 
-
 # Run the Gradio app
 if __name__ == "__main__":
     launch_gradio()