Update app.py

app.py

@@ -1,10 +1,13 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import HTMLResponse
+from fastapi.staticfiles import StaticFiles
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-
 import os
 
 
+
 # Set cache directory for Hugging Face Transformers
 os.environ["TRANSFORMERS_CACHE"] = "/home/user/.cache"
 
@@ -15,15 +18,26 @@ model = AutoModelForCausalLM.from_pretrained("matsant01/STEMerald-2b")
 # Initialize FastAPI app
 app = FastAPI()
 
-
-
-@app.get("/")
-def read_root():
-    return {"message": "Welcome to the STEMerald-2b API"}
-
-#@app.post("/generate/")
-#def generate_text(prompt: str):
-#    inputs = tokenizer(prompt, return_tensors="pt")
-#    outputs = model.generate(inputs["input_ids"], max_length=50)
-#    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-#    return {"generated_text": generated_text}
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+# Serve the HTML file
+@app.get("/", response_class=HTMLResponse)
+async def read_root():
+    with open("index.html", "r") as f:
+        return f.read()
+
+@app.post("/generate/")
+async def generate_text(prompt: str):
+    if not prompt:
+        raise HTTPException(status_code=400, detail="Prompt cannot be empty")
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(inputs["input_ids"], max_length=50)
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return {"generated_text": generated_text}