Update app.py

app.py

@@ -9,35 +9,15 @@ model_id = "sshleifer/tiny-gpt2"
  
 #model_id = "remi/bertabs-finetuned-extractive-abstractive-summarization"
 
-# Ensure cache directory exists
-cache_dir = "/app/cache"
-os.makedirs(cache_dir, exist_ok=True)
-
-# 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
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id)
 
 def generate(text: str) -> str:
     inputs = tokenizer(text, return_tensors="pt")
-    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)
+    outputs = model.generate(**inputs, max_new_tokens=50)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 if __name__ == "__main__":
-    # Example usage
-    prompt = "The meaning of life is <mask>."
+    prompt = "The capital of France is"
     result = generate(prompt)
     print(result)