Update app.py

app.py

@@ -1,4 +1,3 @@
-import spaces
 import gradio as gr
 import torch
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
@@ -11,7 +10,11 @@ print("Loading IndicBART tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(model_name, do_lower_case=False, use_fast=False, keep_accents=True)
 
 print("Loading IndicBART model on CPU...")
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="cpu")
+model = AutoModelForSeq2SeqLM.from_pretrained(
+    model_name, 
+    torch_dtype=torch.float32,  # Use float32 for better CPU performance
+    device_map="cpu"
+)
 
 # Language mapping
 LANGUAGE_CODES = {
@@ -29,11 +32,8 @@ LANGUAGE_CODES = {
     "Telugu": "<2te>"
 }
 
-@spaces.GPU(duration=60)
 def generate_response(input_text, source_lang, target_lang, task_type, max_length):
-    """Generate response using IndicBART"""
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model_gpu = model.to(device)
+    """Generate response using IndicBART on CPU"""
     
     # Get language codes
     src_code = LANGUAGE_CODES[source_lang]
@@ -51,43 +51,45 @@ def generate_response(input_text, source_lang, target_lang, task_type, max_lengt
         formatted_input = f"{input_text} </s> {src_code}"
         decoder_start_token = tgt_code
     
-    # Tokenize input
+    # Tokenize input (keep on CPU)
     inputs = tokenizer(formatted_input, return_tensors="pt", padding=True, truncation=True, max_length=512)
-    inputs = {k: v.to(device) for k, v in inputs.items()}
     
     # Get decoder start token id
-    decoder_start_token_id = tokenizer._convert_token_to_id_with_added_voc(decoder_start_token)
+    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
+    # Generate on CPU
     with torch.no_grad():
-        outputs = model_gpu.generate(
+        outputs = model.generate(
             **inputs,
             decoder_start_token_id=decoder_start_token_id,
             max_length=max_length,
-            num_beams=4,
+            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
+            use_cache=True,
+            do_sample=False  # Deterministic for CPU
         )
     
     # Decode output
     generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
     
-    # Move model back to CPU
-    model_gpu.cpu()
-    torch.cuda.empty_cache()
-    
     return generated_text
 
 # Create Gradio interface
-with gr.Blocks(title="IndicBART Multilingual Assistant", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="IndicBART CPU Multilingual Assistant", theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
-    # 🇮🇳 IndicBART Multilingual Assistant
+    # 🇮🇳 IndicBART Multilingual Assistant (CPU Version)
     
     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():
@@ -104,7 +106,9 @@ with gr.Blocks(title="IndicBART Multilingual Assistant", theme=gr.themes.Soft())
                 interactive=False
             )
             
-            generate_btn = gr.Button("Generate", variant="primary", size="lg")
+            with gr.Row():
+                generate_btn = gr.Button("Generate", variant="primary", size="lg")
+                clear_btn = gr.Button("Clear", variant="secondary")
         
         with gr.Column(scale=1):
             task_type = gr.Dropdown(
@@ -126,9 +130,9 @@ with gr.Blocks(title="IndicBART Multilingual Assistant", theme=gr.themes.Soft())
             )
             
             max_length = gr.Slider(
-                minimum=50,
-                maximum=300,
-                value=100,
+                minimum=20,
+                maximum=200,  # Reduced for faster CPU processing
+                value=80,
                 step=10,
                 label="Max Length"
             )
@@ -137,11 +141,11 @@ with gr.Blocks(title="IndicBART Multilingual Assistant", theme=gr.themes.Soft())
     gr.Markdown("### 💡 Try these examples:")
     
     examples = [
-        ["Hello, how are you?", "English", "Hindi", "Translation", 100],
-        ["मैं एक छात्र हूं", "Hindi", "English", "Translation", 100],
-        ["আমি ভাত খাই", "Bengali", "English", "Translation", 100],
-        ["भारत एक", "Hindi", "Hindi", "Text Completion", 150],
-        ["The capital of India", "English", "English", "Text Completion", 100]
+        ["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(
@@ -151,12 +155,21 @@ with gr.Blocks(title="IndicBART Multilingual Assistant", theme=gr.themes.Soft())
         fn=generate_response
     )
     
-    # Connect generate button
+    # Event handlers
+    def clear_fields():
+        return "", ""
+    
+    # Connect buttons
     generate_btn.click(
         generate_response,
         inputs=[input_text, source_lang, target_lang, task_type, max_length],
         outputs=output_text
     )
+    
+    clear_btn.click(
+        clear_fields,
+        outputs=[input_text, output_text]
+    )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)  # Added share=True for easier access