Update app.py

app.py

@@ -1,38 +1,24 @@
-from fastapi import FastAPI, Request
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
-
-app = FastAPI()
-
-model_id = "google/flan-t5-small" # Replace with your model here
-
-
+import os
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+# Example: using Google Flan-T5 small
+model_id = "google/flan-t5-small"
 #"unsloth/mistral-7b-v0.2-bnb-4bit"
 #deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B
 
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    torch_dtype=torch.float16,
-    device_map="auto",
-)
-
-
-@app.post("/generate")
-async def generate(request: Request):
-    data = await request.json()
-    prompt = data.get("prompt", "").strip()
+# Ensure cache directories exist
+os.makedirs("/app/cache", exist_ok=True)
 
-  
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(
-        inputs.input_ids,
-        max_new_tokens=100,
-        use_cache=True,
-    
-        temperature=0.7,
-    )
+# Load tokenizer & model using the custom cache path
+tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir="/app/cache")
+model = AutoModelForSeq2SeqLM.from_pretrained(model_id, cache_dir="/app/cache")
 
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+# Example simple inference
+def generate(text):
+    inputs = tokenizer(text, return_tensors="pt")
+    outputs = model.generate(**inputs)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    return {"output": generated_text}
+if __name__ == "__main__":
+    # Simple CLI test
+    prompt = "Translate English to French: Hello, how are you?"
+    print(generate(prompt))