Update app.py

app.py

@@ -1,35 +1,25 @@
 from fastapi import FastAPI, Request
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
 
 app = FastAPI()
 
-generator = pipeline("text-generation", model="microsoft/phi-2", max_new_tokens=150)
+# Load model and tokenizer
+model_id = "meta-llama/Llama-3.1-8B"
 
-sessions = {}
+tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-@app.get("/")
-def read_root():
-    return {"message": "LLM API running!"}
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto"  # Will auto-detect if CUDA or CPU
+)
+
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=256)
 
 @app.post("/chat")
 async def chat(request: Request):
     data = await request.json()
-    session_id = data.get("session_id", "default")
-    user_input = data.get("message", "")
-    system_prompt = data.get("system_prompt", "You are a helpful assistant.")
-
-    if session_id not in sessions:
-        sessions[session_id] = system_prompt + "\n"
-
-    sessions[session_id] += f"User: {user_input}\nAssistant:"
-
-    output = generator(sessions[session_id], max_new_tokens=150)[0]["generated_text"]
-
-    if "Assistant:" in output:
-        assistant_response = output.split("Assistant:")[-1].strip()
-    else:
-        assistant_response = output.strip()
-
-    sessions[session_id] += f" {assistant_response}\n"
-
-    return {"response": assistant_response}
+    prompt = data.get("prompt", "")
+    output = pipe(prompt)[0]['generated_text']
+    return {"response": output}