Update app.py

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, TextStreamer
 from huggingface_hub import login
 import PyPDF2
 import pandas as pd
@@ -70,7 +70,7 @@ def load_model(hf_token):
         
         login(token=hf_token)
         
-        # Load tokenizer (requires sentencepiece)
+        # Load tokenizer
         tokenizer = AutoTokenizer.from_pretrained(
             MODEL_NAME,
             token=hf_token
@@ -91,14 +91,17 @@ def load_model(hf_token):
         st.error(f"🤖 Model loading failed: {str(e)}")
         return None
 
-# Generation function for translation
+# Generation function for translation with streaming
 def generate_translation(input_text, model, tokenizer):
     try:
         # Tokenize the input (no prompt needed for seq2seq translation models)
         inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512)
         inputs = inputs.to(DEVICE)
         
-        # Generate translation
+        # Set up the streamer for real-time output
+        streamer = TextStreamer(tokenizer, skip_special_tokens=True)
+        
+        # Generate translation with streaming
         model.eval()
         with torch.no_grad():
             outputs = model.generate(
@@ -107,12 +110,15 @@ def generate_translation(input_text, model, tokenizer):
                 max_length=512,
                 num_beams=5,
                 length_penalty=1.0,
-                early_stopping=True
+                early_stopping=True,
+                streamer=streamer,
+                return_dict_in_generate=True,
+                output_scores=True
             )
         
-        # Decode the output
-        translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return translation
+        # Decode the full output for storage and metrics
+        translation = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
+        return translation, streamer
         
     except Exception as e:
         raise Exception(f"Generation error: {str(e)}")
@@ -154,21 +160,30 @@ if prompt := st.chat_input("Enter text to translate into French..."):
     file_context = process_file(uploaded_file)
     input_text = file_context if file_context else prompt
     
-    # Generate translation
+    # Generate translation with streaming
     if model and tokenizer:
         try:
             with st.chat_message("assistant", avatar=BOT_AVATAR):
                 start_time = time.time()
-                translation = generate_translation(input_text, model, tokenizer)
                 
-                # Display the translation
-                st.markdown(translation)
-                st.session_state.messages.append({"role": "assistant", "content": translation})
+                # Create a placeholder for streaming output
+                response_container = st.empty()
+                full_response = ""
+                
+                # Generate translation and stream output
+                translation, streamer = generate_translation(input_text, model, tokenizer)
+                
+                # Streamlit will automatically display the streamed output via the TextStreamer
+                # Collect the full response for metrics and storage
+                full_response = translation
+                
+                # Update the placeholder with the final response
+                response_container.markdown(full_response)
                 
                 # Calculate performance metrics
                 end_time = time.time()
                 input_tokens = len(tokenizer(input_text)["input_ids"])
-                output_tokens = len(tokenizer(translation)["input_ids"])
+                output_tokens = len(tokenizer(full_response)["input_ids"])
                 speed = output_tokens / (end_time - start_time)
                 
                 # Calculate costs (hypothetical pricing model)
@@ -184,6 +199,9 @@ if prompt := st.chat_input("Enter text to translate into French..."):
                     f"💵 Cost (AOA): {total_cost_aoa:.4f}"
                 )
                 
+                # Store the full response in chat history
+                st.session_state.messages.append({"role": "assistant", "content": full_response})
+                
         except Exception as e:
             st.error(f"⚡ Translation error: {str(e)}")
     else: