Empereur / main.py
Empereur-Pirate's picture
Update main.py
8db8221 verified
raw
history blame
2.46 kB
import os
import requests
from fastapi import FastAPI, Request, Depends
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import Optional, Any
from fastapi.responses import HTMLResponse
app = FastAPI()
# Montage du répertoire statique
app.mount("/static", StaticFiles(directory="static"), name="static")
# Route racine pour servir index.html
@app.get("/", include_in_schema=False)
async def root():
with open("static/index.html", "r") as html_file:
content = html_file.read()
return HTMLResponse(content=content)
# Check whether we are executing inside a Hugging Face Space
SPACE_NAME = os.getenv("SPACE_NAME", default=None)
if SPACE_NAME is not None:
print(f"Running inside {SPACE_NAME} Space.")
try:
# Try to auto-login using the Space's environment variables
login(automatically=True)
except Exception as e:
print(f"Failed to auto-login ({str(e)}). Manually check the HF_ACCESS_TOKEN environment variable.")
sys.exit(1)
try:
HUGGINGFACE_TOKEN = os.environ['HF_ACCESS_TOKEN']
except KeyError:
print('The environment variable "HF_ACCESS_TOKEN" is not found. Please configure it correctly in your Space.')
sys.exit(1)
# Set up the API endpoint and headers
model_id = "152334H/miqu-1-70b-sf"
endpoint = f"https://api-inference.huggingface.co/models/{model_id}"
headers = {"Authorization": f"Bearer {HUGGINGFACE_TOKEN}"}
# Helper function to read raw request bodies
async def parse_raw(request: Request):
return await request.body()
# Generate text using the Inference API
def generate_text(input_text: str) -> str:
data = {"inputs": input_text}
response = requests.post(endpoint, json=data, headers=headers)
print(response.json()) # Print the entire response object
try:
return response.json()["generated_text"]
except KeyError as e:
print(f"KeyError: {e}. Check the response object for the correct key.")
return ""
# Route for generating text
@app.post("/generate_text")
async def generate_text_route(data: BaseModel = Depends(parse_raw)):
input_text = data.decode("utf-8")
if not input_text or len(input_text) <= 0:
return JSONResponse({"error": "Empty input received."}, status_code=400)
return {"output": generate_text(input_text)}
# Mount static files
app.mount("/static", StaticFiles(directory="static"), name="static")