Update app.py

app.py

@@ -12,7 +12,7 @@ tokenizer = AutoTokenizer.from_pretrained(model_name, do_lower_case=False, use_f
 print("Loading IndicBART model on CPU...")
 model = AutoModelForSeq2SeqLM.from_pretrained(
     model_name, 
-    torch_dtype=torch.float32,  # Use float32 for better CPU performance
+    torch_dtype=torch.float32,
     device_map="cpu"
 )
 
@@ -35,50 +35,65 @@ LANGUAGE_CODES = {
 def generate_response(input_text, source_lang, target_lang, task_type, max_length):
     """Generate response using IndicBART on CPU"""
     
-    # Get language codes
-    src_code = LANGUAGE_CODES[source_lang]
-    tgt_code = LANGUAGE_CODES[target_lang]
-    
-    # Format input based on task type
-    if task_type == "Translation":
-        formatted_input = f"{input_text} </s> {src_code}"
-        decoder_start_token = tgt_code
-    elif task_type == "Text Completion":
-        # For completion, use target language
-        formatted_input = f"{input_text} </s> {tgt_code}"
-        decoder_start_token = tgt_code
-    else:  # Text Generation
-        formatted_input = f"{input_text} </s> {src_code}"
-        decoder_start_token = tgt_code
-    
-    # Tokenize input (keep on CPU)
-    inputs = tokenizer(formatted_input, return_tensors="pt", padding=True, truncation=True, max_length=512)
-    
-    # Get decoder start token id
+    if not input_text.strip():
+        return "Please enter some text to process."
+    
     try:
-        decoder_start_token_id = tokenizer._convert_token_to_id_with_added_voc(decoder_start_token)
-    except:
-        # Fallback if the method doesn't exist
-        decoder_start_token_id = tokenizer.convert_tokens_to_ids(decoder_start_token)
-    
-    # Generate on CPU
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            decoder_start_token_id=decoder_start_token_id,
-            max_length=max_length,
-            num_beams=2,  # Reduced for faster CPU inference
-            early_stopping=True,
-            pad_token_id=tokenizer.pad_token_id,
-            eos_token_id=tokenizer.eos_token_id,
-            use_cache=True,
-            do_sample=False  # Deterministic for CPU
+        # Get language codes
+        src_code = LANGUAGE_CODES[source_lang]
+        tgt_code = LANGUAGE_CODES[target_lang]
+        
+        # Format input based on task type
+        if task_type == "Translation":
+            formatted_input = f"{input_text} </s> {src_code}"
+            decoder_start_token = tgt_code
+        elif task_type == "Text Completion":
+            formatted_input = f"{input_text} </s> {tgt_code}"
+            decoder_start_token = tgt_code
+        else:  # Text Generation
+            formatted_input = f"{input_text} </s> {src_code}"
+            decoder_start_token = tgt_code
+        
+        # Tokenize input - KEY FIX: Explicitly set return_token_type_ids=False
+        inputs = tokenizer(
+            formatted_input, 
+            return_tensors="pt", 
+            padding=True, 
+            truncation=True, 
+            max_length=512,
+            return_token_type_ids=False  # This prevents the error
         )
-    
-    # Decode output
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
-    
-    return generated_text
+        
+        # Alternative fix: Remove token_type_ids if present
+        if 'token_type_ids' in inputs:
+            del inputs['token_type_ids']
+        
+        # Get decoder start token id
+        try:
+            decoder_start_token_id = tokenizer._convert_token_to_id_with_added_voc(decoder_start_token)
+        except:
+            decoder_start_token_id = tokenizer.convert_tokens_to_ids(decoder_start_token)
+        
+        # Generate on CPU
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                decoder_start_token_id=decoder_start_token_id,
+                max_length=max_length,
+                num_beams=2,
+                early_stopping=True,
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                use_cache=True,
+                do_sample=False
+            )
+        
+        # Decode output
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
+        return generated_text
+        
+    except Exception as e:
+        return f"Error generating response: {str(e)}"
 
 # Create Gradio interface
 with gr.Blocks(title="IndicBART CPU Multilingual Assistant", theme=gr.themes.Soft()) as demo:
@@ -88,8 +103,6 @@ with gr.Blocks(title="IndicBART CPU Multilingual Assistant", theme=gr.themes.Sof
     Experience IndicBART - trained on **11 Indian languages**! Perfect for translation, text completion, and multilingual generation.
     
     **Supported Languages**: Assamese, Bengali, Gujarati, Hindi, Kannada, Malayalam, Marathi, Oriya, Punjabi, Tamil, Telugu, English
-    
-    *Note: Running on CPU - responses may take longer than GPU version.*
     """)
     
     with gr.Row():
@@ -131,34 +144,39 @@ with gr.Blocks(title="IndicBART CPU Multilingual Assistant", theme=gr.themes.Sof
             
             max_length = gr.Slider(
                 minimum=20,
-                maximum=200,  # Reduced for faster CPU processing
+                maximum=200,
                 value=80,
                 step=10,
                 label="Max Length"
             )
     
-    # Examples
+    # Simplified examples to avoid caching issues
     gr.Markdown("### 💡 Try these examples:")
     
-    examples = [
-        ["Hello, how are you?", "English", "Hindi", "Translation", 80],
-        ["मैं एक छात्र हूं", "Hindi", "English", "Translation", 80],
-        ["আমি ভাত খাই", "Bengali", "English", "Translation", 80],
-        ["भारत एक", "Hindi", "Hindi", "Text Completion", 100],
-        ["The capital of India", "English", "English", "Text Completion", 80]
-    ]
-    
-    gr.Examples(
-        examples=examples,
-        inputs=[input_text, source_lang, target_lang, task_type, max_length],
-        outputs=output_text,
-        fn=generate_response
-    )
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("**English to Hindi**")
+            example1_btn = gr.Button("Hello, how are you?")
+        with gr.Column():
+            gr.Markdown("**Hindi to English**") 
+            example2_btn = gr.Button("मैं एक छात्र हूं")
+        with gr.Column():
+            gr.Markdown("**Bengali to English**")
+            example3_btn = gr.Button("আমি ভাত খাই")
     
     # Event handlers
     def clear_fields():
         return "", ""
     
+    def set_example1():
+        return "Hello, how are you?", "English", "Hindi", "Translation"
+    
+    def set_example2():
+        return "मैं एक छात्र हूं", "Hindi", "English", "Translation"
+    
+    def set_example3():
+        return "আমি ভাত খাই", "Bengali", "English", "Translation"
+    
     # Connect buttons
     generate_btn.click(
         generate_response,
@@ -170,9 +188,28 @@ with gr.Blocks(title="IndicBART CPU Multilingual Assistant", theme=gr.themes.Sof
         clear_fields,
         outputs=[input_text, output_text]
     )
+    
+    example1_btn.click(
+        set_example1,
+        outputs=[input_text, source_lang, target_lang, task_type]
+    )
+    
+    example2_btn.click(
+        set_example2,
+        outputs=[input_text, source_lang, target_lang, task_type]
+    )
+    
+    example3_btn.click(
+        set_example3,
+        outputs=[input_text, source_lang, target_lang, task_type]
+    )
+
+# Launch with all fixes applied
 if __name__ == "__main__":
     demo.launch(
         share=True,
-        server_port=7860,
-        show_error=True
+        ssr_mode=False,           # Disable SSR
+        cache_examples=False,     # Disable example caching - KEY FIX
+        show_error=True,
+        enable_queue=False        # Disable queue to avoid startup issues
     )