Update app.py

app.py

@@ -1,26 +1,43 @@
 import os
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, BertForMaskedLM
 
 # Model identifier
 model_id = "sshleifer/tiny-gpt2"
 #"google/flan-t5-small"
 #"unsloth/mistral-7b-v0.2-bnb-4bit"
 #deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B
+ 
+#model_id = "remi/bertabs-finetuned-extractive-abstractive-summarization"
 
-# Ensure cache directory exists and is writable
+# Ensure cache directory exists
 cache_dir = "/app/cache"
 os.makedirs(cache_dir, exist_ok=True)
 
-# Load tokenizer and model, pointing to custom cache
-tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=cache_dir)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_id, cache_dir=cache_dir)
+# Load using appropriate class
+if model_id.startswith("remi/bertabs"):
+    # BERT masked language model for abstractive summarization
+    model = BertForMaskedLM.from_pretrained(model_id, cache_dir=cache_dir)
+    tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=cache_dir)
+else:
+    # Sequence-to-sequence model like Flan-T5
+    tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=cache_dir)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_id, cache_dir=cache_dir)
+
+# Inference function
 
-# Simple inference function
 def generate(text: str) -> str:
     inputs = tokenizer(text, return_tensors="pt")
-    outputs = model.generate(**inputs)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    if hasattr(model, 'generate'):
+        outputs = model.generate(**inputs)
+        return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    else:
+        # For masked LM, demonstrate mask filling
+        from transformers import pipeline
+        fill = pipeline('fill-mask', model=model, tokenizer=tokenizer, cache_dir=cache_dir)
+        return fill(text)
 
 if __name__ == "__main__":
-    prompt = "Translate English to French: Hello, how are you?"
-    print(generate(prompt))
+    # Example usage
+    prompt = "The meaning of life is <mask>."
+    result = generate(prompt)
+    print(result)