Update app.py

app.py

@@ -34,28 +34,28 @@ nltk.download("punkt", quiet=True)
 
 # Zero-Shot Classification Model (Topic Detection)
 ZSC_MODEL_NAME = "facebook/bart-large-mnli"
-zsc_tokenizer = AutoTokenizer.from_pretrained(ZSC_MODEL_NAME)
-zsc_model = AutoModelForSequenceClassification.from_pretrained(ZSC_MODEL_NAME)
+zsc_tokenizer = AutoTokenizer.from_pretrained(ZSC_MODEL_NAME, force_download=True)
+zsc_model = AutoModelForSequenceClassification.from_pretrained(ZSC_MODEL_NAME, force_download=True)
 zero_shot_classifier = pipeline("zero-shot-classification", model=zsc_model, tokenizer=zsc_tokenizer)
 
 # Summarization Model (Chunk-based Summaries)
 SUM_MODEL_NAME = "facebook/bart-large-cnn"
-sum_tokenizer = AutoTokenizer.from_pretrained(SUM_MODEL_NAME)
-sum_model = AutoModelForSeq2SeqLM.from_pretrained(SUM_MODEL_NAME)
+sum_tokenizer = AutoTokenizer.from_pretrained(SUM_MODEL_NAME, force_download=True)
+sum_model = AutoModelForSeq2SeqLM.from_pretrained(SUM_MODEL_NAME, force_download=True)
 summarizer = pipeline("summarization", model=sum_model, tokenizer=sum_tokenizer)
 
 # QA Model (Chunk-based QA)
 QA_MODEL_NAME = "deepset/roberta-base-squad2"
-qa_tokenizer = AutoTokenizer.from_pretrained(QA_MODEL_NAME)
-qa_model = AutoModelForQuestionAnswering.from_pretrained(QA_MODEL_NAME)
+qa_tokenizer = AutoTokenizer.from_pretrained(QA_MODEL_NAME, force_download=True)
+qa_model = AutoModelForQuestionAnswering.from_pretrained(QA_MODEL_NAME, force_download=True)
 qa_pipeline = pipeline("question-answering", model=qa_model, tokenizer=qa_tokenizer)
 
 # Speech-to-Text (STT) with tiny Whisper
 WHISPER_MODEL_NAME = "openai/whisper-tiny"
-whisper_processor = WhisperProcessor.from_pretrained(WHISPER_MODEL_NAME)
-whisper_model = WhisperForConditionalGeneration.from_pretrained(WHISPER_MODEL_NAME)
+whisper_processor = WhisperProcessor.from_pretrained(WHISPER_MODEL_NAME, force_download=True)
+whisper_model = WhisperForConditionalGeneration.from_pretrained(WHISPER_MODEL_NAME, force_download=True)
 
-# For real-time token usage, let's use tiktoken (GPT-3.5 style tokenizer)
+# For real-time token usage, we'll use tiktoken (GPT-3.5 style tokenizer)
 encoding = tiktoken.get_encoding("cl100k_base")
 
 ###############################################################################
@@ -304,17 +304,19 @@ def answer_question(session_id, question):
 def chat(user_input, chat_history, session_id):
     if session_id not in SESSIONS:
         SESSIONS[session_id] = {"files": {}, "chat_history": []}
+    # If the user wants to search for a reference:
     if user_input.lower().startswith("ref:"):
         query = user_input[4:].strip()
         result = find_reference(session_id, query)
-        chat_history.append((user_input, result))
+        chat_history.append({"role": "assistant", "content": result})
         return "", chat_history
+    # Process the question using QA:
     answer = answer_question(session_id, user_input)
     question_tokens = approximate_tokens(user_input)
     answer_tokens = approximate_tokens(answer)
     usage_str = f"Tokens: Q={question_tokens}, A={answer_tokens}, Total={question_tokens + answer_tokens}"
     full_answer = f"{answer}\n\n({usage_str})"
-    chat_history.append((user_input, full_answer))
+    chat_history.append({"role": "assistant", "content": full_answer})
     return "", chat_history
 
 ###############################################################################
@@ -368,27 +370,27 @@ with gr.Blocks() as demo:
         new_session_btn.click(fn=reset_session, outputs=[session_id, new_session_out])
     
     gr.Markdown("### 2. Voice Input (STT Only)")
-    # Removed the 'source' parameter because it is not supported in this version.
     audio_in = gr.Audio(type="filepath", label="Speak your question")
     stt_btn = gr.Button("Transcribe")
     stt_output = gr.Textbox(label="Transcribed Text")
     stt_btn.click(fn=transcribe_audio, inputs=[audio_in], outputs=[stt_output])
     
     gr.Markdown("### 3. Chat / Q&A (Enter text below)")
+    # Set type="messages" for openai-style chat messages
     chatbot = gr.Chatbot(label="Chat History", type="messages")
     user_input = gr.Textbox(label="Your question (or 'ref: <term>' for reference search)", lines=2)
     send_btn = gr.Button("Send")
     
     def user_message(user_msg, history):
-        history = history + [[user_msg, None]]
+        history = history + [{"role": "user", "content": user_msg}]
         return "", history
     send_btn.click(fn=user_message, inputs=[user_input, chatbot], outputs=[user_input, chatbot], queue=False)
     
     def bot_message(history, sid):
-        # Check if history is empty
         if not history:
             return []
-        user_msg = history[-1][0]
+        # The most recent message should be from the user.
+        user_msg = history[-1]["content"]
         _, updated_history = chat(user_msg, history, sid)
         return updated_history
     send_btn.click(fn=bot_message, inputs=[chatbot, session_id], outputs=[chatbot])