Update app.py

app.py

@@ -1,17 +1,19 @@
+import os
 from fastapi import FastAPI
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
 
+os.environ["HF_HOME"] = "/app/cache"
+
+# Load model
 model_name = "mynuddin/chatbot"  
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name).to("cpu")  
+model = AutoModelForCausalLM.from_pretrained(model_name).to("cpu")
 
 app = FastAPI()
 
 @app.post("/generate")
 def generate_text(prompt: str):
     inputs = tokenizer(prompt, return_tensors="pt")
-    with torch.no_grad():
-        output = model.generate(**inputs, max_length=128)
+    output = model.generate(**inputs, max_length=128)
     generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
     return {"generated_query": generated_text}