Create app.py

app.py

@@ -0,0 +1,148 @@
+import os
+import json
+import torch
+from threading import Thread
+from typing import List, Dict, Any, Optional, Union
+
+from fastapi import FastAPI, HTTPException, Request, BackgroundTasks
+from fastapi.responses import StreamingResponse
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
+from transformers import AutoTokenizer, TextIteratorStreamer
+from vllm import LLM, SamplingParams
+
+
+# Initialize FastAPI app
+app = FastAPI(title="GainEnergy/OGAI-24B API")
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Load environment variables
+MODEL_ID = os.environ.get("MODEL_ID", "GainEnergy/OGAI-24B")
+DEFAULT_MAX_LENGTH = int(os.environ.get("DEFAULT_MAX_LENGTH", "2048"))
+DEFAULT_TEMPERATURE = float(os.environ.get("DEFAULT_TEMPERATURE", "0.7"))
+
+# Initialize the model and tokenizer
+try:
+    model = LLM(
+        model=MODEL_ID,
+        trust_remote_code=True,
+        tensor_parallel_size=torch.cuda.device_count(),
+    )
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+    print(f"Model {MODEL_ID} loaded successfully!")
+except Exception as e:
+    print(f"Error loading model: {e}")
+    raise
+
+# Pydantic models for request/response
+class Message(BaseModel):
+    role: str = Field(..., description="The role of the message sender (system, user, assistant)")
+    content: str = Field(..., description="The content of the message")
+
+class GenerationRequest(BaseModel):
+    messages: List[Message] = Field(..., description="List of messages in the conversation")
+    temperature: Optional[float] = Field(DEFAULT_TEMPERATURE, description="Temperature for sampling")
+    max_tokens: Optional[int] = Field(DEFAULT_MAX_LENGTH, description="Maximum number of tokens to generate")
+    top_p: Optional[float] = Field(0.95, description="Top-p sampling parameter")
+    top_k: Optional[int] = Field(50, description="Top-k sampling parameter")
+    stream: Optional[bool] = Field(False, description="Whether to stream the response")
+
+class GenerationResponse(BaseModel):
+    generated_text: str = Field(..., description="Generated text from the model")
+
+# Helper function to format messages for the model
+def format_messages(messages: List[Message]) -> str:
+    """Format a list of messages into a prompt string the model can understand."""
+    formatted_prompt = ""
+    
+    for message in messages:
+        if message.role == "system":
+            formatted_prompt += f"<|system|>\n{message.content}</s>\n"
+        elif message.role == "user":
+            formatted_prompt += f"<|user|>\n{message.content}</s>\n"
+        elif message.role == "assistant":
+            formatted_prompt += f"<|assistant|>\n{message.content}</s>\n"
+    
+    # Add the final assistant token to prompt the model to generate a response
+    formatted_prompt += "<|assistant|>\n"
+    
+    return formatted_prompt
+
+# API endpoints
+@app.get("/")
+async def root():
+    """Root endpoint with basic information."""
+    return {
+        "status": "running",
+        "model": MODEL_ID,
+        "version": "1.0.0"
+    }
+
+@app.post("/generate", response_model=GenerationResponse)
+async def generate(request: GenerationRequest):
+    """Generate text based on the conversation history."""
+    try:
+        prompt = format_messages(request.messages)
+        
+        sampling_params = SamplingParams(
+            temperature=request.temperature,
+            max_tokens=request.max_tokens,
+            top_p=request.top_p, 
+            top_k=request.top_k
+        )
+        
+        outputs = model.generate(prompt, sampling_params)
+        generated_text = outputs[0].outputs[0].text
+        
+        return {"generated_text": generated_text}
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")
+
+@app.post("/generate_stream")
+async def generate_stream(request: GenerationRequest):
+    """Stream generated text based on the conversation history."""
+    if not request.stream:
+        return await generate(request)
+    
+    try:
+        prompt = format_messages(request.messages)
+        
+        sampling_params = SamplingParams(
+            temperature=request.temperature,
+            max_tokens=request.max_tokens,
+            top_p=request.top_p,
+            top_k=request.top_k
+        )
+        
+        async def stream_generator():
+            for output in model.generate(prompt, sampling_params, stream=True):
+                chunk = output.outputs[0].text
+                yield f"data: {json.dumps({'text': chunk})}\n\n"
+            yield "data: [DONE]\n\n"
+        
+        return StreamingResponse(
+            stream_generator(),
+            media_type="text/event-stream"
+        )
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Streaming generation failed: {str(e)}")
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint."""
+    return {"status": "healthy"}
+
+# Run the FastAPI app with uvicorn when this script is executed directly
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)