update app

app.py

@@ -6,8 +6,14 @@ from transformers import (
     AutoTokenizer, AutoModelForCausalLM
 )
 import torch
+import json
+from fastapi import FastAPI, HTTPException, Depends, Header
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+import uvicorn
+from pydantic import BaseModel
+from typing import Optional
 
-# Configuration for multiple models, can add more by extending MODEL_CONFIGS dict
+# Configuration for multiple models
 MODEL_CONFIGS = {
     "gpt2": {
         "type": "causal",
@@ -39,45 +45,71 @@ MODEL_CONFIGS = {
     }
 }
 
-# Environment variables for optional authentication and private model access
+# Environment variables
 HF_TOKEN = os.getenv("HF_TOKEN")
 API_KEY = os.getenv("API_KEY")
 ADMIN_PASSWORD = os.getenv("ADMIN_PASSWORD")
 
-# Global state for caching loaded model and tokenizer
+# Global state for caching
 loaded_model_name = None
 model = None
 tokenizer = None
 
+# Pydantic models for API
+class GenerateRequest(BaseModel):
+    prompt: str
+    model_name: str = "gpt2"
+    max_length: int = 100
+    temperature: float = 0.7
+    top_p: float = 0.9
+    top_k: int = 50
+
+class GenerateResponse(BaseModel):
+    generated_text: str
+    model_used: str
+    status: str = "success"
+
+# Security
+security = HTTPBearer(auto_error=False)
+
 def load_model_and_tokenizer(model_name):
     global loaded_model_name, model, tokenizer
+    
+    if model_name not in MODEL_CONFIGS:
+        raise ValueError(f"Model {model_name} not supported. Available models: {list(MODEL_CONFIGS.keys())}")
+    
     if model_name == loaded_model_name and model is not None and tokenizer is not None:
         return model, tokenizer
     
-    config = MODEL_CONFIGS[model_name]
-    if HF_TOKEN:
-        tokenizer = config["tokenizer_class"].from_pretrained(model_name, use_auth_token=HF_TOKEN)
-        model = config["model_class"].from_pretrained(model_name, use_auth_token=HF_TOKEN)
-    else:
-        tokenizer = config["tokenizer_class"].from_pretrained(model_name)
-        model = config["model_class"].from_pretrained(model_name)
-
-    # Set pad token for causal models if missing (important for generation padding)
-    if config["type"] == "causal" and tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-
-    loaded_model_name = model_name
-    return model, tokenizer
-
-def authenticate_api_key(key):
-    if API_KEY and key != API_KEY:
-        return False
-    return True
+    try:
+        config = MODEL_CONFIGS[model_name]
+        
+        # Load tokenizer and model
+        if HF_TOKEN:
+            tokenizer = config["tokenizer_class"].from_pretrained(model_name, use_auth_token=HF_TOKEN)
+            model = config["model_class"].from_pretrained(model_name, use_auth_token=HF_TOKEN)
+        else:
+            tokenizer = config["tokenizer_class"].from_pretrained(model_name)
+            model = config["model_class"].from_pretrained(model_name)
 
-def generate_text(prompt, model_name, max_length, temperature, top_p, top_k, api_key=""):
-    if API_KEY and not authenticate_api_key(api_key):
-        return "Error: Invalid API key"
+        # Set pad token for causal models if missing
+        if config["type"] == "causal" and tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        loaded_model_name = model_name
+        return model, tokenizer
+        
+    except Exception as e:
+        raise RuntimeError(f"Failed to load model {model_name}: {str(e)}")
 
+def authenticate_api_key(credentials: Optional[HTTPAuthorizationCredentials] = Depends(security)):
+    if API_KEY:
+        if not credentials or credentials.credentials != API_KEY:
+            raise HTTPException(status_code=401, detail="Invalid or missing API key")
+    return True
+
+def generate_text_core(prompt, model_name, max_length, temperature, top_p, top_k):
+    """Core text generation function"""
     try:
         config = MODEL_CONFIGS[model_name]
         model, tokenizer = load_model_and_tokenizer(model_name)
@@ -96,11 +128,9 @@ def generate_text(prompt, model_name, max_length, temperature, top_p, top_k, api
                     num_return_sequences=1
                 )
             generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-            # Return generated continuation (remove original prompt)
             return generated_text[len(prompt):].strip()
 
         elif config["type"] == "seq2seq":
-            # Add task prefix for certain seq2seq models like flan-t5
             task_prompt = f"Complete this text: {prompt}" if "flan-t5" in model_name.lower() else prompt
             inputs = tokenizer(task_prompt, return_tensors="pt", max_length=512, truncation=True)
             with torch.no_grad():
@@ -117,8 +147,62 @@ def generate_text(prompt, model_name, max_length, temperature, top_p, top_k, api
             return generated_text.strip()
 
     except Exception as e:
-        return f"Error generating text: {str(e)}"
+        raise RuntimeError(f"Error generating text: {str(e)}")
 
+# Gradio interface function
+def generate_text_gradio(prompt, model_name, max_length, temperature, top_p, top_k, api_key=""):
+    if API_KEY and api_key != API_KEY:
+        return "Error: Invalid API key"
+    
+    try:
+        return generate_text_core(prompt, model_name, max_length, temperature, top_p, top_k)
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Create FastAPI app
+app = FastAPI(title="Multi-Model Text Generation API", version="1.0.0")
+
+# API Routes
+@app.post("/generate", response_model=GenerateResponse)
+async def generate_text_api(
+    request: GenerateRequest,
+    authenticated: bool = Depends(authenticate_api_key)
+):
+    try:
+        generated_text = generate_text_core(
+            request.prompt,
+            request.model_name,
+            request.max_length,
+            request.temperature,
+            request.top_p,
+            request.top_k
+        )
+        return GenerateResponse(
+            generated_text=generated_text,
+            model_used=request.model_name
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/models")
+async def list_models():
+    return {
+        "models": [
+            {
+                "name": name,
+                "description": config["description"],
+                "size": config["size"],
+                "type": config["type"]
+            }
+            for name, config in MODEL_CONFIGS.items()
+        ]
+    }
+
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy", "loaded_model": loaded_model_name}
+
+# Create Gradio interface
 with gr.Blocks(title="Multi-Model Text Generation Server") as demo:
     gr.Markdown("# Multi-Model Text Generation Server")
     gr.Markdown("Choose a model from the dropdown, enter a text prompt, and generate text.")
@@ -131,6 +215,16 @@ with gr.Blocks(title="Multi-Model Text Generation Server") as demo:
                 value="gpt2",
                 interactive=True
             )
+            
+            # Show model info
+            model_info = gr.Markdown("**Model Info:** Original GPT-2, good for creative writing (117M)")
+            
+            def update_model_info(model_name):
+                config = MODEL_CONFIGS[model_name]
+                return f"**Model Info:** {config['description']} ({config['size']})"
+            
+            model_selector.change(update_model_info, inputs=model_selector, outputs=model_info)
+            
             prompt_input = gr.Textbox(
                 label="Text Prompt",
                 placeholder="Enter the text prompt here...",
@@ -171,7 +265,7 @@ with gr.Blocks(title="Multi-Model Text Generation Server") as demo:
             )
 
     generate_btn.click(
-        fn=generate_text,
+        fn=generate_text_gradio,
         inputs=[prompt_input, model_selector, max_length_slider, temperature_slider, top_p_slider, top_k_slider, api_key_input],
         outputs=output_textbox
     )
@@ -186,13 +280,61 @@ with gr.Blocks(title="Multi-Model Text Generation Server") as demo:
         inputs=prompt_input
     )
 
-auth_config = ("admin", ADMIN_PASSWORD) if ADMIN_PASSWORD else None
+    # API documentation
+    with gr.Accordion("API Documentation", open=False):
+        gr.Markdown("""
+        ## REST API Endpoints
+        
+        ### POST /generate
+        Generate text using the specified model.
+        
+        **Request Body:**
+        ```json
+        {
+            "prompt": "Your text prompt here",
+            "model_name": "gpt2",
+            "max_length": 100,
+            "temperature": 0.7,
+            "top_p": 0.9,
+            "top_k": 50
+        }
+        ```
+        
+        **Response:**
+        ```json
+        {
+            "generated_text": "Generated text...",
+            "model_used": "gpt2",
+            "status": "success"
+        }
+        ```
+        
+        ### GET /models
+        List all available models.
+        
+        ### GET /health
+        Check server health and loaded model status.
+        
+        **Example cURL:**
+        ```bash
+        curl -X POST "http://localhost:7860/generate" \
+             -H "Content-Type: application/json" \
+             -H "Authorization: Bearer YOUR_API_KEY" \
+             -d '{"prompt": "Once upon a time", "model_name": "gpt2"}'
+        ```
+        """)
+
+# Mount Gradio app to FastAPI
+app = gr.mount_gradio_app(app, demo, path="/")
 
 if __name__ == "__main__":
+    auth_config = ("admin", ADMIN_PASSWORD) if ADMIN_PASSWORD else None
+    
+    # Launch with both FastAPI and Gradio
     demo.launch(
         auth=auth_config,
-    #    share=True,        # Required for Spaces if localhost isn't accessible
         server_name="0.0.0.0",
         server_port=7860,
-        ssr_mode=False     # Optional: disable server-side rendering to avoid Svelte i18n error
-    )
+        ssr_mode=False,
+        share=False
+    )

requirements.txt

@@ -1,4 +1,7 @@
-gradio>=3.50.0
-transformers>=4.30.0
-torch>=2.0.0
-tokenizers>=0.13.0
+gradio>=4.0.0
+transformers>=4.21.0
+torch>=1.12.0
+fastapi>=0.68.0
+uvicorn>=0.15.0
+pydantic>=1.8.0
+python-multipart>=0.0.5