Update app.py

app.py

@@ -1,7 +1,18 @@
 from fastapi import FastAPI
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
 app = FastAPI()
 
-@app.get("/")
-def greet_json():
-    return {"Hello": "World!"}
+tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/deepseek-llm-7b-base")
+model = AutoModelForCausalLM.from_pretrained(
+    "deepseek-ai/deepseek-llm-7b-base",
+    torch_dtype=torch.float16,
+    device_map="auto"
+)
+
+@app.get("/ask")
+def ask(prompt: str):
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    outputs = model.generate(**inputs, max_new_tokens=100)
+    return {"response": tokenizer.decode(outputs[0], skip_special_tokens=True)}