Update app.py

app.py

@@ -11,6 +11,8 @@ import os
 from transformers import pipeline
 from transformers.pipelines.audio_utils import ffmpeg_read
 
+from sentence_transformers import SentenceTransformer
+
 from langchain.prompts import PromptTemplate
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
@@ -19,10 +21,20 @@ from langchain_text_splitters import SentenceTransformersTokenTextSplitter
 
 from PIL import Image
 
-from transformers import LlavaNextProcessor, LlavaNextForConditionalGeneration, TextIteratorStreamer
-processor = LlavaNextProcessor.from_pretrained("llava-hf/llava-v1.6-mistral-7b-hf")
-model = LlavaNextForConditionalGeneration.from_pretrained("llava-hf/llava-v1.6-mistral-7b-hf", torch_dtype=torch.float16, low_cpu_mem_usage=True)
-model.to("cuda:0")
+# from transformers import LlavaNextProcessor, LlavaNextForConditionalGeneration, TextIteratorStreamer
+# processor = LlavaNextProcessor.from_pretrained("llava-hf/llava-v1.6-mistral-7b-hf")
+# model = LlavaNextForConditionalGeneration.from_pretrained("llava-hf/llava-v1.6-mistral-7b-hf", torch_dtype=torch.float16, low_cpu_mem_usage=True)
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="auto",
+    ).to("cuda:0")
+terminators = [
+    tokenizer.eos_token_id,
+    tokenizer.convert_tokens_to_ids("<|eot_id|>")
+    ]
 
 embeddings_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
 
@@ -35,6 +47,7 @@ from huggingface_hub import InferenceClient
 """
 For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
 """
+# client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
 device = 0 if torch.cuda.is_available() else "cpu"
@@ -47,7 +60,7 @@ asr_pl = pipeline(
 )
 
 application_title = "Enlight Innovations Limited -- Demo"
-application_description = "This demo is desgined to illustrate our basic ideas and feasibility in implementation."
+application_description = "This demo is designed to illustrate our basic ideas and feasibility in implementation."
 
 @spaces.GPU
 def respond(
@@ -69,8 +82,38 @@ def respond(
 
     messages.append({"role": "user", "content": message})
 
-    response = ""
+    input_ids = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        return_tensors="pt"
+    ).to(model.device)
+
+    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
+
+    generate_kwargs = dict(
+        input_ids= input_ids,
+        streamer=streamer,
+        max_new_tokens=max_new_tokens,
+        do_sample=True,
+        temperature=temperature,
+        eos_token_id=terminators,
+    )
+    # This will enforce greedy generation (do_sample=False) when the temperature is passed 0, avoiding the crash.             
+    if temperature == 0:
+        generate_kwargs['do_sample'] = False
+        
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+
+    outputs = []
+    for text in streamer:
+        outputs.append(text)
+        #print(outputs)
+        yield "".join(outputs)
 
+"""
+    response = ""
+        
     for message in client.chat_completion(
         messages,
         max_tokens=max_tokens,
@@ -82,6 +125,7 @@ def respond(
 
         response += token
         yield response
+"""
 
 @spaces.GPU
 def transcribe(asr_inputs, task):
@@ -90,7 +134,7 @@ def transcribe(asr_inputs, task):
         raise gr.Error("No audio file submitted! Please upload or record an audio file before submitting your request.") 
         
     text = asr_pl(asr_inputs, batch_size=ASR_BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)["text"]
-    return text
+    return text.strip()
 
 
 """Gradio User Interface"""