Update app.py

app.py

@@ -6,7 +6,6 @@ import torch
 import logging
 import traceback
 import sys
-import ast
 
 logging.basicConfig(
     level=logging.INFO,
@@ -122,6 +121,7 @@ def process_and_summarize(audio_file, translate, model_size, do_summarize=True):
         raise gr.Error(f"Processing failed: {str(e)}")
 
 
+
 @spaces.GPU(duration=60)
 def answer_question(context, question):
     logger.info("Starting Q&A process")
@@ -135,25 +135,19 @@ def answer_question(context, question):
             {"role": "user", "content": f"Context: {context}\n\nQuestion: {question}"},
         ]
         
-        outputs = qa_pipeline(messages, max_new_tokens=256)
+        out = qa_pipeline(messages, max_new_tokens=256)
         
-        logger.info(f"Raw model output: {outputs}")
-
-        # Parse the string output into a list of dictionaries
-        try:
-            parsed_output = ast.literal_eval(outputs)
-            if isinstance(parsed_output, list):
-                for item in parsed_output:
-                    if isinstance(item, dict) and item.get('role') == 'assistant':
-                        answer = item.get('content', '')
-                        break
-                else:
-                    answer = "No assistant response found in the model's output."
-            else:
-                answer = str(parsed_output)
-        except (ValueError, SyntaxError):
-            # If parsing fails, return the raw output
-            answer = str(outputs)
+        logger.info(f"Raw model output: {out}")
+
+        generated_text = out[0]['generated_text']
+        
+        # Find the assistant's message
+        for message in generated_text:
+            if message['role'] == 'assistant':
+                answer = message['content']
+                break
+        else:
+            answer = "No assistant response found in the model's output."
 
         logger.info(f"Extracted answer: {answer}")
         return answer
@@ -162,6 +156,7 @@ def answer_question(context, question):
         logger.error(traceback.format_exc())
         return f"Error occurred during Q&A process. Please try again. Error: {str(e)}"
 
+
 # Main interface
 with gr.Blocks() as iface:
     gr.Markdown("# WhisperX Audio Transcription, Translation, Summarization, and Q&A (with ZeroGPU support)")