Update app.py

app.py

@@ -1,20 +1,19 @@
 import gradio as gr
 from audio_processing import process_audio, print_results
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForQuestionAnswering
 import spaces
 import torch
 
 # Check if CUDA is available
 cuda_available = torch.cuda.is_available()
+device = "cuda" if cuda_available else "cpu"
 
 # Initialize the summarization and question-answering models
-summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
-qa_model = pipeline("question-answering", model="distilbert-base-cased-distilled-squad")
+summarizer_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn").to(device)
+summarizer_tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
 
-# Move models to GPU if available
-if cuda_available:
-    summarizer.to('cuda')
-    qa_model.to('cuda')
+qa_model = AutoModelForQuestionAnswering.from_pretrained("distilbert-base-cased-distilled-squad").to(device)
+qa_tokenizer = AutoTokenizer.from_pretrained("distilbert-base-cased-distilled-squad")
 
 @spaces.GPU
 def transcribe_audio(audio_file, translate, model_size):
@@ -41,13 +40,19 @@ def transcribe_audio(audio_file, translate, model_size):
 
 @spaces.GPU
 def summarize_text(text):
-    summary = summarizer(text, max_length=150, min_length=50, do_sample=False)[0]['summary_text']
+    inputs = summarizer_tokenizer(text, max_length=1024, truncation=True, return_tensors="pt").to(device)
+    summary_ids = summarizer_model.generate(inputs["input_ids"], max_length=150, min_length=50, do_sample=False)
+    summary = summarizer_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
     return summary
 
 @spaces.GPU
 def answer_question(context, question):
-    result = qa_model(question=question, context=context)
-    return result['answer']
+    inputs = qa_tokenizer(question, context, return_tensors="pt").to(device)
+    outputs = qa_model(**inputs)
+    answer_start = torch.argmax(outputs.start_logits)
+    answer_end = torch.argmax(outputs.end_logits) + 1
+    answer = qa_tokenizer.decode(inputs["input_ids"][0][answer_start:answer_end])
+    return answer
 
 @spaces.GPU
 def process_and_summarize(audio_file, translate, model_size):