Update app.py

app.py

@@ -1,39 +1,24 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from fastapi import FastAPI
 
-# Model ID on Hugging Face
-MODEL_ID = "rasyosef/Phi-1_5-Instruct-v0.1"
 
-# Load tokenizer and model from local cache (pre-downloaded in Docker build)
+MODEL_ID = "MBZUAI/LaMini-Flan-T5-77M"
+
+
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
+model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID)
 
 app = FastAPI()
 
-@app.get("/")
-def root():
-    return {
-        "message": "✅ Trigger82 AI API is running!",
-        "usage": "Use /chat?query=Your+question to get a response."
-    }
-
 @app.get("/chat")
 def chat(query: str):
     """
     GET /chat?query=Your+question
     Returns JSON: {"answer": "...model’s reply..."}
     """
-    # Build the instruction‐style prompt expected by Phi‐1.5 Instruct
-    prompt = (
-        "<|im_start|>system\nYou are a helpful assistant.<|im_end|>"
-        "<|im_start|>user\n" + query + "<|im_end|>"
-        "<|im_start|>assistant\n"
-    )
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=200)
-    # Only decode newly generated tokens (skip the “prompt” tokens)
+    inputs = tokenizer(query, return_tensors="pt")
+    outputs = model.generate(**inputs, max_new_tokens=100)
     response = tokenizer.decode(
-        outputs[0][inputs.input_ids.shape[-1]:], 
-        skip_special_tokens=True
+        outputs[0], skip_special_tokens=True
     )
     return {"answer": response.strip()}