Update app.py

app.py

@@ -1,9 +1,9 @@
 from fastapi import FastAPI
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import uvicorn
 
 tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
-model = AutoModel.from_pretrained("Qwen/Qwen3-0.6B")
+model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B")
 
 app = FastAPI()
 
@@ -14,11 +14,7 @@ def greet_json():
 @app.get("/message")
 async def message(input: str):
     inputs = tokenizer(input, return_tensors="pt", padding=True, truncation=True)
-    output = model.generate(
-        input_ids=inputs["input_ids"],
-        attention_mask=inputs["attention_mask"],
-        temperature=0.3
-    )
+    output = model.generate(**inputs, max_length=50, temperature=0.3)
     return tokenizer.decode(output[0], skip_special_tokens=True)
 
 if __name__ == "__main__":