Update streamlit_app.py

streamlit_app.py

@@ -1,9 +1,10 @@
 # streamlit_app.py
-# A minimal Streamlit app rebuilt with the LangChain framework.
+# A robust Streamlit app with proper error handling and fallback options
 
 import streamlit as st
 import torch
-from transformers import pipeline
+import logging
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 
 # Updated LangChain imports for modern versions
 from langchain_community.llms import HuggingFacePipeline
@@ -11,122 +12,326 @@ from langchain.prompts import PromptTemplate
 from langchain.chains import LLMChain
 from langchain.memory import ConversationBufferMemory
 
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 # -----------------------------------------------------------------------------
-# CORE MODEL LOGIC (Rebuilt with LangChain)
+# CORE MODEL LOGIC (Rebuilt with LangChain and Error Handling)
 # -----------------------------------------------------------------------------
 class LangChainBot:
     def __init__(self):
         """
-        Loads the models and wraps them in LangChain components.
+        Loads the models and wraps them in LangChain components with fallback options.
         """
+        self.chain = None
+        self.translator = None
+        self.memory = None
+        
         try:
-            # 1. Load the base Hugging Face pipelines
-            generator_pipeline = pipeline(
-                "text2text-generation",
-                model="ai4bharat/IndicBARTSS",
-                device=0 if torch.cuda.is_available() else -1,
-                torch_dtype=(torch.float16 if torch.cuda.is_available() else torch.float32),
-                max_new_tokens=150,
-                repetition_penalty=1.2
-            )
+            # Check CUDA availability
+            device = 0 if torch.cuda.is_available() else -1
+            torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+            
+            st.info(f"Using device: {'CUDA' if device == 0 else 'CPU'}")
+            
+            # Try to load the main model with error handling
+            self._load_main_model(device, torch_dtype)
             
-            # Added `trust_remote_code=True` to allow the special translator model to load.
-            self.translator = pipeline(
-                "translation",
-                model="ai4bharat/indictrans2-indic-indic-1B",
-                device=0 if torch.cuda.is_available() else -1,
-                trust_remote_code=True
-            )
-
-            # 2. Wrap the generator in a LangChain LLM object
-            llm = HuggingFacePipeline(pipeline=generator_pipeline)
-
-            # 3. Create a Prompt Template
-            template = """
-            You are a helpful conversational AI. Respond to the user's message.
-
-            {history}
-            मनुष्य: {input}
-            सहायक:
-            """
-            prompt_template = PromptTemplate(input_variables=["history", "input"], template=template)
-
-            # 4. Set up conversational memory
-            self.memory = ConversationBufferMemory(memory_key="history")
-
-            # 5. Create the final LLMChain
-            self.chain = LLMChain(
-                llm=llm,
-                prompt=prompt_template,
-                verbose=True,
-                memory=self.memory
-            )
+            # Try to load the translator with error handling
+            self._load_translator(device)
             
         except Exception as e:
-            st.error(f"Fatal: Could not load models. Error: {e}")
-            self.chain = None
-            self.translator = None
+            logger.error(f"Fatal error during initialization: {e}")
+            st.error(f"Fatal: Could not initialize the bot. Error: {e}")
+
+    def _load_main_model(self, device, torch_dtype):
+        """Load the main generation model with fallback options."""
+        models_to_try = [
+            "ai4bharat/IndicBARTSS",
+            "google/flan-t5-small",  # Fallback option
+            "t5-small"  # Another fallback
+        ]
+        
+        for model_name in models_to_try:
+            try:
+                st.info(f"Attempting to load model: {model_name}")
+                
+                # Try loading with pipeline first
+                generator_pipeline = pipeline(
+                    "text2text-generation",
+                    model=model_name,
+                    device=device,
+                    torch_dtype=torch_dtype,
+                    max_new_tokens=150,
+                    repetition_penalty=1.2,
+                    trust_remote_code=True  # Added this for compatibility
+                )
+                
+                # Wrap in LangChain LLM
+                llm = HuggingFacePipeline(pipeline=generator_pipeline)
+                
+                # Create prompt template
+                template = """
+                You are a helpful conversational AI. Respond to the user's message appropriately.
+                Previous conversation:
+                {history}
+                
+                Human: {input}
+                Assistant:
+                """
+                prompt_template = PromptTemplate(
+                    input_variables=["history", "input"], 
+                    template=template
+                )
+                
+                # Set up memory
+                self.memory = ConversationBufferMemory(memory_key="history")
+                
+                # Create the chain
+                self.chain = LLMChain(
+                    llm=llm,
+                    prompt=prompt_template,
+                    verbose=True,
+                    memory=self.memory
+                )
+                
+                st.success(f"Successfully loaded model: {model_name}")
+                return  # Success, exit the loop
+                
+            except Exception as e:
+                logger.warning(f"Failed to load {model_name}: {e}")
+                st.warning(f"Failed to load {model_name}, trying next option...")
+                continue
+        
+        raise Exception("All model loading attempts failed")
+
+    def _load_translator(self, device):
+        """Load the translator with fallback options."""
+        translators_to_try = [
+            "ai4bharat/indictrans2-indic-indic-1B",
+            "Helsinki-NLP/opus-mt-en-hi",  # Fallback for English-Hindi
+        ]
+        
+        for translator_name in translators_to_try:
+            try:
+                st.info(f"Attempting to load translator: {translator_name}")
+                
+                self.translator = pipeline(
+                    "translation",
+                    model=translator_name,
+                    device=device,
+                    trust_remote_code=True
+                )
+                
+                st.success(f"Successfully loaded translator: {translator_name}")
+                return  # Success
+                
+            except Exception as e:
+                logger.warning(f"Failed to load translator {translator_name}: {e}")
+                st.warning(f"Failed to load translator {translator_name}, trying next option...")
+                continue
+        
+        st.warning("No translator loaded - translation features will be limited")
 
     def _translate(self, text, source_lang, target_lang):
-        """Translation logic remains the same."""
+        """Translation logic with improved error handling."""
         if not self.translator or source_lang == target_lang:
             return text
+            
         try:
-            codes = {'english': 'eng_Latn', 'hindi': 'hin_Deva', 'tamil': 'tam_Taml', 'telugu': 'tel_Telu'}
-            result = self.translator(text, src_lang=codes[source_lang], tgt_lang=codes[target_lang])
-            return result[0]['translation_text']
+            # Define language codes
+            codes = {
+                'english': 'eng_Latn', 
+                'hindi': 'hin_Deva', 
+                'tamil': 'tam_Taml', 
+                'telugu': 'tel_Telu'
+            }
+            
+            if source_lang in codes and target_lang in codes:
+                result = self.translator(
+                    text, 
+                    src_lang=codes[source_lang], 
+                    tgt_lang=codes[target_lang]
+                )
+                return result[0]['translation_text']
+            else:
+                # Fallback for simple English-Hindi translation
+                if source_lang == 'english' and target_lang == 'hindi':
+                    result = self.translator(text)
+                    return result[0]['translation_text'] if result else text
+                
         except Exception as e:
-            st.warning(f"Translation failed. Error: {e}")
-            return text
+            logger.warning(f"Translation failed: {e}")
+            st.warning(f"Translation failed, using original text. Error: {e}")
+            
+        return text
 
     def get_response(self, user_message, input_lang, output_lang):
-        """The main function to get a response."""
+        """Generate response with comprehensive error handling."""
         if not self.chain:
-            return "Error: The LangChain chain is not initialized."
+            return "Error: The LangChain chain is not initialized. Please check the logs above."
 
-        hindi_message = self._translate(user_message, input_lang, 'hindi')
-        hindi_response = self.chain.run(hindi_message)
-        final_response = self._translate(hindi_response, 'hindi', output_lang)
-
-        return final_response
+        try:
+            # Translate input to a common language if needed
+            if input_lang != 'english':
+                processed_message = self._translate(user_message, input_lang, 'english')
+            else:
+                processed_message = user_message
+            
+            # Generate response
+            response = self.chain.run(processed_message)
+            
+            # Translate output if needed
+            if output_lang != 'english':
+                final_response = self._translate(response, 'english', output_lang)
+            else:
+                final_response = response
+                
+            return final_response
+            
+        except Exception as e:
+            logger.error(f"Error generating response: {e}")
+            return f"I apologize, but I encountered an error while processing your request: {str(e)}"
 
 # -----------------------------------------------------------------------------
-# MINIMAL STREAMLIT UI (This part remains mostly the same)
+# STREAMLIT UI WITH BETTER ERROR HANDLING
 # -----------------------------------------------------------------------------
 
-st.set_page_config(layout="centered")
-st.title("LangChain Model Interface")
+st.set_page_config(
+    page_title="LangChain Model Interface",
+    page_icon="🤖",
+    layout="centered"
+)
+
+st.title("🤖 LangChain Model Interface")
+st.markdown("*Multi-language conversational AI powered by LangChain*")
 
+# Initialize the bot with progress tracking
 @st.cache_resource
 def load_bot():
-    return LangChainBot()
+    with st.spinner("Loading models... This may take a few minutes on first run."):
+        return LangChainBot()
 
+# Load the bot
 bot = load_bot()
 
-if bot and bot.chain: # Only show the UI if the bot loaded successfully
+# Check if bot loaded successfully
+if bot and bot.chain:
+    st.success("✅ Bot loaded successfully!")
+    
     st.markdown("---")
+    
+    # Language selection
     language_options = ["english", "hindi", "tamil", "telugu"]
     col1, col2 = st.columns(2)
+    
     with col1:
-        input_lang = st.selectbox("Input Language", options=language_options, index=0)
+        input_lang = st.selectbox(
+            "🔤 Input Language", 
+            options=language_options, 
+            index=0,
+            help="Select the language you'll type in"
+        )
     with col2:
-        output_lang = st.selectbox("Output Language", options=language_options, index=1)
+        output_lang = st.selectbox(
+            "🗣️ Output Language", 
+            options=language_options, 
+            index=1,
+            help="Select the language for the response"
+        )
 
-    user_input = st.text_area("Your Message:", height=100)
+    # Chat interface
+    st.markdown("### 💬 Chat Interface")
+    user_input = st.text_area(
+        "Your Message:", 
+        height=100,
+        placeholder=f"Type your message in {input_lang}..."
+    )
 
-    if st.button("Get Response"):
-        if user_input:
-            with st.spinner("LangChain is processing your request..."):
-                response = bot.get_response(user_input, input_lang, output_lang)
-                st.markdown("### Model Response:")
+    col1, col2 = st.columns([3, 1])
+    
+    with col1:
+        if st.button("🚀 Get Response", type="primary"):
+            if user_input.strip():
+                with st.spinner("🤔 LangChain is processing your request..."):
+                    response = bot.get_response(user_input, input_lang, output_lang)
+                    
+                st.markdown("### 🤖 Model Response:")
                 st.info(response)
-        else:
-            st.warning("Please enter a message.")
-
-    # Add a button to clear LangChain's memory
-    if st.button("Clear Conversation Memory"):
-        if hasattr(bot, 'memory'):
-            bot.memory.clear()
-            st.success("Conversation memory has been cleared.")
+                
+                # Add to conversation history display
+                if 'conversation_history' not in st.session_state:
+                    st.session_state.conversation_history = []
+                
+                st.session_state.conversation_history.append({
+                    'user': user_input,
+                    'bot': response,
+                    'input_lang': input_lang,
+                    'output_lang': output_lang
+                })
+                
+            else:
+                st.warning("⚠️ Please enter a message.")
+    
+    with col2:
+        if st.button("🧹 Clear Memory"):
+            if hasattr(bot, 'memory') and bot.memory:
+                bot.memory.clear()
+                if 'conversation_history' in st.session_state:
+                    del st.session_state.conversation_history
+                st.success("✅ Conversation memory cleared!")
+
+    # Display conversation history
+    if 'conversation_history' in st.session_state and st.session_state.conversation_history:
+        st.markdown("### 📝 Conversation History")
+        for i, conv in enumerate(reversed(st.session_state.conversation_history[-5:])):  # Show last 5
+            with st.expander(f"Exchange {len(st.session_state.conversation_history) - i}"):
+                st.markdown(f"**You ({conv['input_lang']})**: {conv['user']}")
+                st.markdown(f"**Bot ({conv['output_lang']})**: {conv['bot']}")
+
 else:
-    st.error("Application could not start. Please check the logs.")
+    st.error("❌ Application could not start. Please check the error messages above.")
+    
+    # Show some troubleshooting tips
+    st.markdown("### 🔧 Troubleshooting Tips:")
+    st.markdown("""
+    1. **Model Loading Issues**: The models might be too large for the available resources
+    2. **Memory Issues**: Try restarting the application
+    3. **Network Issues**: Ensure stable internet connection for model downloads
+    4. **Compatibility Issues**: Some models might not be compatible with the current environment
+    """)
+    
+    if st.button("🔄 Retry Loading"):
+        st.cache_resource.clear()
+        st.rerun()
+
+# Add sidebar with information
+with st.sidebar:
+    st.markdown("### ℹ️ Information")
+    st.markdown("""
+    This application uses:
+    - **LangChain** for conversation management
+    - **Hugging Face Transformers** for AI models
+    - **Multi-language support** via translation models
+    
+    **Supported Languages:**
+    - English
+    - Hindi  
+    - Tamil
+    - Telugu
+    """)
+    
+    if torch.cuda.is_available():
+        st.success("🚀 CUDA GPU detected - faster processing!")
+    else:
+        st.info("💻 Using CPU - processing may be slower")
+    
+    st.markdown("### 🔧 System Status")
+    st.markdown(f"- PyTorch: {torch.__version__}")
+    st.markdown(f"- Device: {'CUDA' if torch.cuda.is_available() else 'CPU'}")
+    if bot and bot.chain:
+        st.markdown("- Model: ✅ Loaded")
+        st.markdown(f"- Translator: {'✅ Loaded' if bot.translator else '❌ Not loaded'}")
+    else:
+        st.markdown("- Model: ❌ Failed to load")