Update app.py

app.py

@@ -1,16 +1,19 @@
-from transformers import pipeline
-from fastapi import FastAPI, Request, HTTPException
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import torch
 
-app = FastAPI()
-    
-generator = pipeline('text2text-generation', model='google/flan-t5-small')
-    
-@app.post("/")
-async def generate(request: Request):
-    try:
-        data = await request.json()
-        text = data["text"]
-        output = generator(text)
-        return {"output": output[0]["generated_text"]}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+model_id = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
+
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16,
+    device_map="auto",
+    trust_remote_code=True
+)
+
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+def generate(prompt):
+    return pipe(prompt, max_new_tokens=100)[0]["generated_text"]
+
+# Launch your app (Gradio, FastAPI, etc.)