Add simple Tailwind UI and serve via FastAPI static files

main.py

@@ -4,6 +4,7 @@ import torch
 import gc
 from fastapi import FastAPI, HTTPException
 from fastapi.staticfiles import StaticFiles
+from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from transformers import pipeline
 
@@ -12,8 +13,17 @@ logger = logging.getLogger(__name__)
 
 app = FastAPI(title="LaMini-LM API", description="API for text generation using LaMini-GPT-774M", version="1.0.0")
 
-# Mount static files
-app.mount("/", StaticFiles(directory="static", html=True), name="static")
+# Add CORS middleware to allow UI requests
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Adjust for production to specific origins
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Mount static files at /ui
+app.mount("/ui", StaticFiles(directory="static", html=True), name="static")
 
 class TextGenerationRequest(BaseModel):
     instruction: str
@@ -28,7 +38,7 @@ def load_model():
     if generator is None:
         try:
             logger.info("Loading LaMini-GPT-774M model...")
-            generator = pipeline('text-generation', model='MBZUAI/LaMini-GPT-774M', device=-1)
+            generator = pipeline('text-generation', model='MBZUAI/LaMini-GPT-774M', device=-1, trust_remote_code=True)
             logger.info("Model loaded successfully.")
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
@@ -36,14 +46,19 @@ def load_model():
         except Exception as e:
             logger.error(f"Failed to load model: {str(e)}")
             generator = None
-            raise Exception(f"Model loading failed: {str(e)}")
+            raise HTTPException(status_code=503, detail=f"Model loading failed: {str(e)}")
 
 @app.get("/health")
 async def health_check():
     return {"status": "healthy"}
 
+@app.get("/")
+async def root():
+    return {"message": "Welcome to the LaMini-LM API. Use POST /generate to generate text or visit /ui for the web interface."}
+
 @app.post("/generate")
 async def generate_text(request: TextGenerationRequest):
+    logger.info(f"Received request: {request.dict()}")
     if generator is None:
         load_model()
     if generator is None:
@@ -67,7 +82,8 @@ async def generate_text(request: TextGenerationRequest):
             temperature=request.temperature,
             top_p=request.top_p,
             num_return_sequences=1,
-            do_sample=True
+            do_sample=True,
+            truncation=True
         )
         generated_text = outputs[0]['generated_text'].replace(wrapper, "").strip()
         return {"generated_text": generated_text}

static/index.html

@@ -126,9 +126,13 @@
       }
 
       try {
-        const response = await fetch('/generate', {
+        const response = await fetch('https://usmansafder-lamini-lm-api.hf.space/generate', {
           method: 'POST',
-          headers: { 'Content-Type': 'application/json' },
+          headers: {
+            'Content-Type': 'application/json',
+            // Add Authorization header if Space is private
+            // 'Authorization': 'Bearer <your_hf_token>'
+          },
           body: JSON.stringify({
             instruction,
             max_length: maxLength,
@@ -141,7 +145,7 @@
           resultDiv.classList.remove('hidden');
           generatedText.textContent = data.generated_text;
         } else {
-          showError(data.detail?.[0]?.msg || 'Failed to generate text.');
+          showError(data.detail?.[0]?.msg || data.detail || 'Failed to generate text.');
         }
       } catch (err) {
         showError('Error connecting to the API. Please try again.');