Update app.py with the latest changes

app.py

@@ -1,31 +1,25 @@
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-import logging
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger("LLM-API")
+# Load the model and tokenizer
+MODEL_NAME = "deepseek-ai/DeepSeek-V3-Base"  # Change to the model you want
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto").to(device)
 
 app = FastAPI()
 
-# Define the input schema
-class InputText(BaseModel):
+class Query(BaseModel):
     input_text: str
 
-@app.get("/")
-def root():
-    logger.info("Root endpoint called.")
-    return {"message": "Welcome to the LLM API"}
-
 @app.post("/predict")
-async def predict(data: InputText, request: Request):
-    logger.info("Received request: %s", await request.body())
-    try:
-        # Log the received input
-        input_text = data.input_text
-        logger.info(f"Processing input: {input_text}")
-        # Return a mock response for now
-        return {"response": f"The input was: {input_text}"}
-    except Exception as e:
-        logger.error(f"Error occurred: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
+async def predict(query: Query):
+    input_text = query.input_text
+    if not input_text:
+        raise HTTPException(status_code=400, detail="Input text cannot be empty.")
+    inputs = tokenizer(input_text, return_tensors="pt").to(device)
+    outputs = model.generate(inputs["input_ids"], max_new_tokens=50, temperature=0.7)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return {"response": response}