Create app.py

app.py

@@ -0,0 +1,112 @@
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import torch
+import threading
+import time
+
+app = FastAPI()
+
+# Global variables to store the model and tokenizer
+model = None
+tokenizer = None
+model_loading_lock = threading.Lock()
+model_loaded = False  # Status flag to indicate if the model is loaded
+
+def load_model(model_name="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"):
+    global model, tokenizer, model_loaded
+    with model_loading_lock:
+        if not model_loaded:
+            print("Loading model...")
+            tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                device_map="sequential",
+                torch_dtype=torch.float16,
+                trust_remote_code=True,
+                low_cpu_mem_usage=True,
+                offload_folder="offload"
+            )
+            model_loaded = True
+            print("Model loaded successfully.")
+        else:
+            print("Model already loaded.")
+
+def check_model_status():
+    """Check if the model is loaded and reload if necessary."""
+    global model_loaded
+    if not model_loaded:
+        print("Model not loaded. Reloading...")
+        load_model()
+    return model_loaded
+
+@app.post("/chat")
+async def chat_endpoint(message: str, temperature: float = 0.7, max_new_tokens: int = 2048):
+    global model, tokenizer
+    
+    # Ensure the model is loaded before proceeding
+    if not check_model_status():
+        raise HTTPException(status_code=503, detail="Model is not ready. Please try again later.")
+    
+    stop_tokens = ["|im_end|"]
+    prompt = f"Human: {message}\n\nAssistant:"
+    
+    # Tokenize the input
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    
+    # Stream the response
+    start_time = time.time()
+    token_count = 0
+    
+    # Create a TextStreamer for token streaming
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    
+    generate_kwargs = dict(
+        input_ids=inputs.input_ids,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id,
+        streamer=streamer  # Use the TextStreamer here
+    )
+    
+    # Start generation in a separate thread
+    threading.Thread(target=model.generate, kwargs=generate_kwargs).start()
+    
+    def generate_response():
+        outputs = []
+        for new_token in streamer:
+            outputs.append(new_token)
+            token_count += 1
+            
+            # Calculate tokens per second
+            elapsed_time = time.time() - start_time
+            tokens_per_second = token_count / elapsed_time if elapsed_time > 0 else 0
+            
+            # Yield the current output and token status
+            yield f"data: {new_token}\n\n"
+            
+            if any(stop_token in new_token for stop_token in stop_tokens):
+                break
+    
+    return StreamingResponse(generate_response(), media_type="text/event-stream")
+
+@app.post("/reload-model")
+async def reload_model():
+    """Reload the model manually via an API endpoint."""
+    global model_loaded
+    model_loaded = False
+    load_model()
+    return {"message": "Model reloaded successfully."}
+
+@app.get("/status")
+async def get_model_status():
+    """Check the status of the model."""
+    status = "Model is loaded and ready." if model_loaded else "Model is not loaded."
+    return {"status": status}
+
+# Load the model when the server starts
+if __name__ == "__main__":
+    load_model()  # Pre-load the model
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)