Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from audio_processing import process_audio
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
 import spaces
 import torch
 import logging
@@ -16,8 +16,18 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+def load_translation_model() :
+    tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
+    model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
+    return model, tokenizer
 
 
+def alternate_translation(inputs): 
+    model, tokenizer = load_translation_model()
+    translated_tokens = model.generate(**inputs, forced_bos_token_id=tokenizer.convert_tokens_to_ids("en"), max_length=100)
+    return tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]    
+    
+    
 def load_qa_model():
     logger.info("Loading Q&A model...")
     try:
@@ -134,7 +144,25 @@ def answer_question(context, question):
             {"role": "system", "content": "You are a helpful assistant who can answer questions based on the given context."},
             {"role": "user", "content": f"Context: {context}\n\nQuestion: {question}"},
         ]
+        alternate_system_message = """
+        You are an AI assistant designed to analyze speech transcriptions in a safe and responsible manner. 
+        Your purpose is to assist people, not to monitor or detect threats.
+
+        When responding to user queries, your primary goals are:
         
+        1. To provide factual, accurate information to the best of your abilities.
+        2. To guide users towards appropriate resources and authorities if they are facing an emergency or urgent situation.
+        3. To refrain from speculating about or escalating potentially concerning situations without clear justification.
+        4. To avoid making judgements or taking actions that could infringe on individual privacy or civil liberties.
+        
+        However, if the speech suggests someone may be in immediate danger or that a crime is being planned, you should:
+
+        - Identify & report 
+        - Identify any cryptic information and report it. 
+        - Avoid probing for additional details or speculating about the nature of the potential threat.
+        - Do not provide any information that could enable or encourage harmful, illegal or unethical acts.
+        Your role is to be a helpful, informative assistant. 
+        """
         out = qa_pipeline(messages, max_new_tokens=256)
         
         logger.info(f"Raw model output: {out}")
@@ -174,6 +202,9 @@ with gr.Blocks() as iface:
     answer_button = gr.Button("Get Answer")
     answer_output = gr.Textbox(label="Answer")
 
+    translate_alternate = gr.Button("Alternate Translation")
+    translate_alternate_output = gr.Textbox(label="Alternate Translation")
+    
     def update_summarize_checkbox(translate):
         return gr.Checkbox(interactive=translate)
 
@@ -191,6 +222,12 @@ with gr.Blocks() as iface:
         outputs=[answer_output]
     )
 
+    translate_alternate.click(
+        alternate_translation, 
+        inputs=[trancription_output],
+        outputs=[translate_alternate_output]
+    )
+    
     gr.Markdown(
         f"""
         ## System Information