Update app.py

app.py

@@ -1,8 +1,8 @@
 import spaces
 import torch
 import gradio as gr
-from transformers import AutoTokenizer, LlamaForCausalLM
-import bitsandbytes, flash_attn
+from transformers import pipeline
+from llama_cpp import Llama
 import os
 
 MODEL_NAME = "openai/whisper-large-v3-turbo"
@@ -19,16 +19,8 @@ pipe = pipeline(
     device=device,
 )
 
-# Load tokenizer and model for SOAP note generation
-tokenizer = AutoTokenizer.from_pretrained("NousResearch/Hermes-3-Llama-3.1-8B", trust_remote_code=True)
-model = LlamaForCausalLM.from_pretrained(
-    "NousResearch/Hermes-3-Llama-3.1-8B",
-    torch_dtype=torch.float16,
-    device_map="auto",
-    load_in_8bit=False,
-    load_in_4bit=True,
-    use_flash_attention_2=True
-)
+# Load the Llama model for SOAP note generation
+llm = Llama(model_path="model.gguf", n_ctx=8000, n_threads=2, chat_format="chatml")
 
 # Prompt for SOAP note generation
 sys_prompt = "You are a world class clinical assistant."
@@ -51,12 +43,17 @@ def transcribe(inputs, task):
     text = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)["text"]
     return text
 
-# Function to generate SOAP notes using LLM
+# Function to generate SOAP notes using Llama model
 def generate_soap(transcribed_text):
-    prompt = f"<|im_start|>system\n{sys_prompt}<|im_end|>\n<|im_start|>user\n{task_prompt}\n{transcribed_text}<|im_end|>\n<|im_start|>assistant"
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cuda")
-    generated_ids = model.generate(input_ids, max_new_tokens=2048, temperature=0.8, repetition_penalty=1.1, do_sample=True, eos_token_id=tokenizer.eos_token_id)
-    response = tokenizer.decode(generated_ids[0][input_ids.shape[-1]:], skip_special_tokens=True, clean_up_tokenization_space=True)
+    prompt = [{"role": "system", "content": sys_prompt}]
+    prompt.append({"role": "user", "content": f"{task_prompt}\n{transcribed_text}"})
+    
+    # Generate a response using the Llama model in streaming mode
+    stream_response = llm.create_chat_completion(messages=prompt, temperature=0.7, max_tokens=2048, stream=True)
+    response = ""
+    for chunk in stream_response:
+        if "content" in chunk['choices'][0]["delta"]:
+            response += chunk['choices'][0]["delta"]["content"]
     return response
 
 # Gradio Interfaces for different inputs