try catch

app.py

@@ -52,27 +52,41 @@ def load_llamaindex_stack(model_id: str, max_new_tokens: int, temperature: float
     tok = AutoTokenizer.from_pretrained(model_id)
     mdl = AutoModelForSeq2SeqLM.from_pretrained(model_id)
     text2text = pipeline(
-        task="text2text-generation",
+        "text2text-generation",
         model=mdl,
         tokenizer=tok,
         max_new_tokens=max_new_tokens,
         temperature=float(temperature)
     )
     """
+    #llm = HuggingFaceLLM(pipeline=text2text)
 
     # Wrap the same tiny HF model for LlamaIndex
     
-    llm = HuggingFaceLLM(
+    config = AutoConfig.from_pretrained(model_id)
+    if config.model_type in ["t5", "mt5", "bart", "mbart", "pegasus", "marian", "prophetnet"]:
+        task = "text2text-generation"   # encoder-decoder / seq2seq
+    else:
+        task = "text-generation" 
+
+     try:
+        llm = HuggingFaceLLM(
         model_name=model_id,
-        tokenizer_name=model_id,
-        model_cls=AutoModelForSeq2SeqLM,
+        tokenizer_name=model_id,        
+        task=task,    
         context_window=2048,
         generate_kwargs={"max_new_tokens": max_new_tokens, "temperature": temperature},
         device_map="cpu",
-    )
-    
-    #llm = HuggingFaceLLM(pipeline=text2text)
-
+        )
+    except  TypeError:
+         llm = HuggingFaceLLM(
+            model_name=model_id,
+            tokenizer_name=model_id,
+            context_window=2048,
+            generate_kwargs={"max_new_tokens": max_new_tokens, "temperature": float(temperature)},
+            device_map="cpu",
+        )
+  
     Settings.embed_model = embed
     Settings.llm = llm