Simplify Q&A and improve UI layout

.gitignore

@@ -1,3 +1,4 @@
 *.pyc
 __pycache__/
 .env
+gradiouitest.py

app.py

@@ -28,7 +28,7 @@ model = SALM.from_pretrained("nvidia/canary-qwen-2.5b").bfloat16().eval().to(dev
 @spaces.GPU
 def transcribe_audio(audio_filepath):
     if audio_filepath is None:
-        return "Please upload an audio file", "", [], 0
+        return "Please upload an audio file", "", []
 
     # Load audio with torchaudio (handles all formats)
     audio, sample_rate = torchaudio.load(audio_filepath)
@@ -56,47 +56,41 @@ def transcribe_audio(audio_filepath):
 
     # Convert output IDs to text
     transcript = model.tokenizer.ids_to_text(output_ids[0].cpu())
-    initial_message = [{"role": "assistant", "content": f"Transcript ready. Ask me questions about it."}]
-    return transcript, transcript, initial_message, 0
+    initial_message = [{"role": "assistant", "content": "Transcript ready. Ask me questions about it."}]
+    return transcript, transcript, initial_message
 
 
-# Enhanced Q&A function with conversation history
+# Simple Q&A function
 @spaces.GPU
-def answer_question_with_history(transcript, question, history, qa_count):
+def answer_question(transcript, question, history):
     if not transcript:
-        return history, qa_count, "Please transcribe audio first"
+        return history, "Please transcribe audio first"
     
-    if qa_count >= 5:
-        history.append({"role": "user", "content": question})
-        history.append({"role": "assistant", "content": "You've reached the maximum of 5 questions for this transcript. Please transcribe new audio to continue."})
-        return history, qa_count, ""
+    if not question:
+        return history, ""
     
-    # Build context from history for better continuity
-    context = ""
-    for msg in history[-4:]:  # Use last 2 exchanges for context
-        if msg.get("role") == "user":
-            context += f"Previous question: {msg['content']}\n"
-    
-    with torch.inference_mode(), model.llm.disable_adapter():
-        prompt = f"{context}Current question: {question}\n\nTranscript:\n{transcript}"
+    with torch.inference_mode():
+        prompt = f"Based on this transcript, answer the following question:\n\nTranscript: {transcript}\n\nQuestion: {question}\n\nAnswer:"
         output_ids = model.generate(
             prompts=[[{"role": "user", "content": prompt}]],
-            max_new_tokens=400,
+            max_new_tokens=256,
         )
 
-    # Convert output IDs to text and extract answer
-    answer = model.tokenizer.ids_to_text(output_ids[0].cpu())
-    answer = answer.split("<|im_start|>assistant")[-1]
+    # Convert output IDs to text
+    full_response = model.tokenizer.ids_to_text(output_ids[0].cpu())
     
-    # Add follow-up prompt if under 5 questions
-    if qa_count < 4:
-        answer += f"\n\nQuestion {qa_count + 1}/5 - What else would you like to know?"
+    # Extract just the answer part
+    if "Answer:" in full_response:
+        answer = full_response.split("Answer:")[-1].strip()
     else:
-        answer += "\n\nThis is your final question for this transcript."
+        answer = full_response.strip()
+    
+    # Clean up any remaining tags
+    answer = answer.replace("<|im_end|>", "").replace("<|im_start|>", "").strip()
     
     history.append({"role": "user", "content": question})
     history.append({"role": "assistant", "content": answer})
-    return history, qa_count + 1, ""
+    return history, ""
 
 # Build the Gradio interface
 with gr.Blocks(theme=theme) as demo:
@@ -105,9 +99,7 @@ with gr.Blocks(theme=theme) as demo:
 
     # State variables
     transcript_state = gr.State()
-    qa_history = gr.State([])
-    qa_counter = gr.State(0)
-
+    
     with gr.Row():
         with gr.Column(scale=1):
             gr.Markdown("### Audio Input")
@@ -130,20 +122,20 @@ with gr.Blocks(theme=theme) as demo:
     gr.Markdown("### Interactive Q&A")
     chatbot = gr.Chatbot(
         type="messages",
-        height=400,
-        label="Conversation History",
+        height=450,
+        label="",
         bubble_full_width=False
     )
     
     with gr.Row():
         question_input = gr.Textbox(
-            label="Your Question",
-            placeholder="e.g., What was the main topic? Why did they say that?",
-            scale=4
+            label="",
+            placeholder="Ask a question about the transcript...",
+            scale=5,
+            container=False
         )
-        ask_btn = gr.Button("Ask", variant="primary", scale=1)
-    
-    clear_chat_btn = gr.Button("Clear Chat", variant="secondary")
+        ask_btn = gr.Button("Ask", variant="primary", scale=1, size="lg")
+        clear_chat_btn = gr.Button("Clear", variant="secondary", scale=1, size="lg")
 
     gr.Markdown("""
     ### Example Questions to Try:
@@ -158,25 +150,25 @@ with gr.Blocks(theme=theme) as demo:
     transcribe_btn.click(
         fn=transcribe_audio,
         inputs=[audio_input],
-        outputs=[transcript_output, transcript_state, chatbot, qa_counter]
+        outputs=[transcript_output, transcript_state, chatbot]
     )
     
     ask_btn.click(
-        fn=answer_question_with_history,
-        inputs=[transcript_state, question_input, chatbot, qa_counter],
-        outputs=[chatbot, qa_counter, question_input]
+        fn=answer_question,
+        inputs=[transcript_state, question_input, chatbot],
+        outputs=[chatbot, question_input]
     )
     
     question_input.submit(
-        fn=answer_question_with_history,
-        inputs=[transcript_state, question_input, chatbot, qa_counter],
-        outputs=[chatbot, qa_counter, question_input]
+        fn=answer_question,
+        inputs=[transcript_state, question_input, chatbot],
+        outputs=[chatbot, question_input]
     )
     
     clear_chat_btn.click(
-        fn=lambda t: ([{"role": "assistant", "content": "Chat cleared. Ask me questions about the transcript."}] if t else [], 1 if t else 0),
-        inputs=[transcript_state],
-        outputs=[chatbot, qa_counter]
+        fn=lambda: [],
+        inputs=[],
+        outputs=[chatbot]
     )
 
 demo.queue()