Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import torch
 from transformers import pipeline
-from ctransformers import AutoModelForCausalLM
+from ctransformers import AutoModelForCausalLM, AutoTokenizer
 
 MODEL_NAME = "openai/whisper-tiny"
 BATCH_SIZE = 8
@@ -15,17 +15,20 @@ pipe = pipeline(
     device=device,
 )
 
+# Set gpu_layers to the number of layers to offload to GPU. Set to 0 if no GPU acceleration is available on your system.
+llm = AutoModelForCausalLM.from_pretrained("TheBloke/Mistral-7B-v0.1-GGUF", model_file="mistral-7b-v0.1.Q4_K_M.gguf", model_type="mistral", gpu_layers=0)
+tokenizer = AutoTokenizer.from_pretrained(llm)
+
+llm_pipe = pipeline("text-generation", model=llm, tokenizer=tokenizer)
+
 def transcribe(inputs, task = "transcribe"):
     if inputs is None:
         raise gr.Error("No audio file submitted! Please upload or record an audio file before submitting your request.")
 
     text = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)["text"]
-    return  text
+    return  llm_pipe(text, max_new_tokens=256)
 
-# Set gpu_layers to the number of layers to offload to GPU. Set to 0 if no GPU acceleration is available on your system.
-llm = AutoModelForCausalLM.from_pretrained("TheBloke/Mistral-7B-v0.1-GGUF", model_file="mistral-7b-v0.1.Q4_K_M.gguf", model_type="mistral", gpu_layers=0)
 
-print(llm("AI is going to"))
 
 iface = gr.Interface(
     fn=transcribe,