added chat endpoint

main.py

@@ -6,7 +6,7 @@ from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from dotenv import load_dotenv
 from accelerate import Accelerator
-
+from typing import List, Tuple
 # Load environment variables from a .env file (useful for local development)
 load_dotenv()
 
@@ -59,45 +59,43 @@ app = FastAPI(
     docs_url="/",  # URL for Swagger docs
     redoc_url="/doc"  # URL for ReDoc docs
 )
-# Set your Hugging Face token from environment variable
-HF_TOKEN = os.getenv("HF_TOKEN")
 
+HF_TOKEN = os.getenv("HF_TOKEN")
 MODEL = "meta-llama/Llama-3.2-1B-Instruct"
-
-# Auto-select CPU or GPU
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device: {device}")
 
-# Set PyTorch to use all available CPU cores if running on CPU
 torch.set_num_threads(multiprocessing.cpu_count())
-
-# Initialize Accelerator for managing device allocation
 accelerator = Accelerator()
 
-# Load model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL, token=HF_TOKEN, use_fast=True)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL,
     token=HF_TOKEN,
     torch_dtype=torch.float16,
-    device_map=device,
-    low_cpu_mem_usage=True,
-
+    device_map=device
 )
 
-# Prepare model for multi-device setup with accelerate
 model, tokenizer = accelerator.prepare(model, tokenizer)
-
-# Pydantic model for input
+# Pydantic models for request validation
 class PromptRequest(BaseModel):
     prompt: str
     max_new_tokens: int = 100
     temperature: float = 0.7
 
+class ChatRequest(BaseModel):
+    message: str
+    history: List[Tuple[str, str]] = []
+    max_new_tokens: int = 100
+    temperature: float = 0.7
+    system_prompt: str = "You are a helpful assistant."
+
+
+# Endpoints
 @app.post("/generate/")
 async def generate_text(request: PromptRequest):
     inputs = tokenizer(request.prompt, return_tensors="pt").to(device)
-    
+   
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
@@ -106,6 +104,36 @@ async def generate_text(request: PromptRequest):
             do_sample=False,
             pad_token_id=tokenizer.eos_token_id
         )
-    
+   
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return {"response": response}
+
+@app.post("/chat/")
+async def chat(request: ChatRequest):
+    conversation = [
+        {"role": "system", "content": request.system_prompt}
+    ]
+    for human, assistant in request.history:
+        conversation.extend([
+            {"role": "user", "content": human},
+            {"role": "assistant", "content": assistant}
+        ])
+    conversation.append({"role": "user", "content": request.message})
+
+    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt").to(device)
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            input_ids,
+            max_new_tokens=request.max_new_tokens,
+            temperature=request.temperature,
+            do_sample=False,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Extract only the assistant's response
+    assistant_response = response.split("Assistant:")[-1].strip()
+    
+    return {"response": assistant_response}