Update app.py

app.py

@@ -6,8 +6,7 @@ from pydantic import BaseModel, field_validator
 from transformers import (
     AutoConfig,
     pipeline,
-    AutoModelForSeq2SeqLM,
-    AutoModelForCausalLM,
+    AutoModel,
     AutoTokenizer,
     GenerationConfig,
     StoppingCriteriaList
@@ -83,18 +82,10 @@ class S3ModelLoader:
                 s3_uri, local_files_only=True
             )
             
-            if "llama" in model_name:
-                model = AutoModelForCausalLM.from_pretrained(
-                s3_uri, config=config, local_files_only=True, rope_scaling = {"type": "linear", "factor": 2.0}
-            )
-            elif 'qwen' in model_name:
-                model = AutoModelForCausalLM.from_pretrained(
-                s3_uri, config=config, local_files_only=True
-            )
-            else:
-                model = AutoModelForSeq2SeqLM.from_pretrained(
+            model = AutoModel.from_pretrained(
                 s3_uri, config=config, local_files_only=True
             )
+            
                 
             tokenizer = AutoTokenizer.from_pretrained(
                 s3_uri, config=config, local_files_only=True
@@ -115,19 +106,11 @@ class S3ModelLoader:
                     model_name, config=config, token=HUGGINGFACE_HUB_TOKEN
                 )
                 
-                if "llama" in model_name:
-                    model = AutoModelForCausalLM.from_pretrained(
-                    model_name, config=config, token=HUGGINGFACE_HUB_TOKEN, rope_scaling = {"type": "linear", "factor": 2.0}
-                )
-                elif 'qwen' in model_name:
-                    model = AutoModelForCausalLM.from_pretrained(
-                    model_name, config=config, token=HUGGINGFACE_HUB_TOKEN
-                )
-                else:
-                    model = AutoModelForSeq2SeqLM.from_pretrained(
+                model = AutoModel.from_pretrained(
                     model_name, config=config, token=HUGGINGFACE_HUB_TOKEN
                 )
 
+
                 if tokenizer.eos_token_id is not None and \
                    tokenizer.pad_token_id is None:
                     tokenizer.pad_token_id = config.pad_token_id \
@@ -164,23 +147,26 @@ async def generate(request: GenerateRequest):
             load_model_and_tokenizer(model_name)
         device = "cuda" if torch.cuda.is_available() else "cpu"
         model.to(device)
+        
+        if "text-to-text" == task_type:
+            generation_config = GenerationConfig(
+                temperature=temperature,
+                max_new_tokens=max_new_tokens,
+                top_p=top_p,
+                top_k=top_k,
+                repetition_penalty=repetition_penalty,
+                do_sample=do_sample,
+                num_return_sequences=num_return_sequences,
+            )
 
-        generation_config = GenerationConfig(
-            temperature=temperature,
-            max_new_tokens=max_new_tokens,
-            top_p=top_p,
-            top_k=top_k,
-            repetition_penalty=repetition_penalty,
-            do_sample=do_sample,
-            num_return_sequences=num_return_sequences,
-        )
-
-        return StreamingResponse(
-            stream_text(model, tokenizer, input_text,
-                         generation_config, stop_sequences,
-                         device, chunk_delay),
-            media_type="text/plain"
-        )
+            return StreamingResponse(
+                stream_text(model, tokenizer, input_text,
+                             generation_config, stop_sequences,
+                             device, chunk_delay),
+                media_type="text/plain"
+            )
+        else:
+            return HTTPException(status_code=400, detail="Task type not text-to-text")
 
     except Exception as e:
         raise HTTPException(