Update app.py

app.py

@@ -1,26 +1,58 @@
 import gradio as gr
+from fastapi import FastAPI
+from pydantic import BaseModel, Field
+from typing import Optional
 
-# Load the model directly using Gradio's load method
-interface = gr.load("models/meta-llama/Llama-3.2-3B-Instruct")
+# Initialize FastAPI app
+app = FastAPI()
 
-# Define the function to call the model with additional parameters
-def custom_chat_fn(prompt, max_length=100, temperature=0.7):
-    # Call the Gradio-loaded model with prompt and parameters
-    response = interface(prompt, max_length=max_length, temperature=temperature)
-    return response[0]['generated_text']
+# Load the model once at startup
+model = gr.load("models/meta-llama/Llama-3.2-3B-Instruct")
 
-# Create a new interface with parameter inputs using the updated Gradio components
-parameterized_interface = gr.Interface(
-    fn=custom_chat_fn,
-    inputs=[
-        gr.Textbox(label="Prompt"),
-        gr.Slider(50, 500, step=10, label="Max Length", value=100),
-        gr.Slider(0.1, 1.0, step=0.1, label="Temperature", value=0.7)
-    ],
-    outputs="text",
-    live=True,
-    allow_flagging='never',
-    description="Enter a prompt, set the max length and temperature, then click 'Submit' to generate text."
-)
+class PoemRequest(BaseModel):
+    prompt: str = Field(..., description="The prompt for poem generation")
+    temperature: Optional[float] = Field(0.7, ge=0.1, le=2.0, description="Controls randomness in generation")
+    top_p: Optional[float] = Field(0.9, ge=0.1, le=1.0, description="Nucleus sampling parameter")
+    top_k: Optional[int] = Field(50, ge=1, le=100, description="Top-k sampling parameter")
+    max_length: Optional[int] = Field(200, ge=50, le=500, description="Maximum length of generated text")
+    repetition_penalty: Optional[float] = Field(1.1, ge=1.0, le=2.0, description="Penalty for repetition")
 
-parameterized_interface.launch()
+class PoemResponse(BaseModel):
+    poem: str
+    parameters_used: dict
+
+@app.post("/generate_poem")
+async def generate_poem(request: PoemRequest) -> PoemResponse:
+    """
+    Generate a poem based on the provided prompt and parameters.
+    
+    Returns:
+        PoemResponse: Contains the generated poem and the parameters used
+    """
+    try:
+        # Prepare generation parameters
+        generation_config = {
+            "temperature": request.temperature,
+            "top_p": request.top_p,
+            "top_k": request.top_k,
+            "max_length": request.max_length,
+            "repetition_penalty": request.repetition_penalty,
+        }
+        
+        # Generate the poem
+        response = model(
+            request.prompt,
+            **generation_config
+        )
+        
+        return PoemResponse(
+            poem=response,
+            parameters_used=generation_config
+        )
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)