Update app.py

app.py

@@ -1,22 +1,44 @@
-#from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from fastapi import FastAPI, Request
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
 app = FastAPI()
-model_id = "google/flan-t5-small"
+
+model_id = "google/flan-t5-small" # Replace with your model here
+
 
 #"unsloth/mistral-7b-v0.2-bnb-4bit"
 #deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B
 
-
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
-generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16,
+    device_map="auto",
+)
+
+cache = {}
 
 @app.post("/generate")
 async def generate(request: Request):
     data = await request.json()
-    prompt = data.get("prompt", "")
-    result = generator(prompt, max_new_tokens=100)[0]["generated_text"]
-    return {"output": result}
+    prompt = data.get("prompt", "").strip()
+
+    if prompt in cache:
+        return {"output": cache[prompt], "cached": True}
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(
+        inputs.input_ids,
+        max_new_tokens=100,
+        use_cache=True,
+        do_sample=True,
+        top_p=0.95,
+        top_k=50,
+        temperature=0.7,
+    )
+
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    cache[prompt] = generated_text
 
+    return {"output": generated_text, "cached": False}