Update app.py

app.py

@@ -1,17 +1,27 @@
+import gradio as gr
+import numpy as np
+import random
+import torch
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
-import random
-import numpy as np
-
-app = FastAPI()
+import uvicorn
+import threading
 
 # Configuration
 model_id = "raduqus/reco_1b_16bit"
 device = "cuda"
 MAX_SEED = np.iinfo(np.int32).max
 
+# Load model globally
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id, 
+    torch_dtype=torch.float16, 
+    variant="fp16", 
+    use_safetensors=True
+).to(device)
+
 def infer(
     prompt, 
     negative_prompt=None, 
@@ -28,18 +38,6 @@ def infer(
     # Set random generator
     generator = torch.Generator().manual_seed(seed)
     
-    # Load model
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id, 
-        torch_dtype=torch.float16, 
-        variant="fp16", 
-        use_safetensors=True
-    )
-    
-    # Move to GPU
-    model = model.to(device)
-    
     # Generate
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
     outputs = model.generate(
@@ -56,6 +54,9 @@ def infer(
     
     return generated_text
 
+# FastAPI for API endpoint
+api = FastAPI()
+
 class RecommendationRequest(BaseModel):
     prompt: str
     negative_prompt: str = None
@@ -65,7 +66,7 @@ class RecommendationRequest(BaseModel):
     temperature: float = 0.7
     top_p: float = 0.9
 
-@app.post("/recommend")
+@api.post("/recommend")
 async def recommend_task(request: RecommendationRequest):
     try:
         result = infer(
@@ -81,6 +82,28 @@ async def recommend_task(request: RecommendationRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+# Gradio Interface
+def gradio_infer(prompt):
+    return infer(prompt)
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=gradio_infer,
+    inputs=gr.Textbox(label="Prompt"),
+    outputs=gr.Textbox(label="Recommendation"),
+    title="Task Recommender",
+    description="Generate task recommendations"
+)
+
+# Function to start FastAPI server
+def start_api_server():
+    uvicorn.run(api, host="0.0.0.0", port=7860)
+
+# Main execution
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    # Start API server in a separate thread
+    api_thread = threading.Thread(target=start_api_server)
+    api_thread.start()
+    
+    # Launch Gradio interface
+    iface.launch()