helperAi / app /main.py
sanmmarr29's picture
Update app/main.py
13993db verified
from fastapi import FastAPI, Request
from fastapi.templating import Jinja2Templates
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse
from transformers import pipeline
from .config import settings
from pydantic import BaseModel
app = FastAPI(
title="DeepSeek Chat",
description="A chat API using DeepSeek model",
version="1.0.0"
)
# Mount static files and templates
app.mount("/static", StaticFiles(directory="app/static"), name="static")
templates = Jinja2Templates(directory="app/templates")
# Initialize pipeline
print("Loading model pipeline...")
pipe = pipeline(
"text-generation",
model=settings.MODEL_NAME,
token=settings.HUGGINGFACE_TOKEN,
trust_remote_code=True
)
class ChatMessage(BaseModel):
message: str
@app.get("/", response_class=HTMLResponse)
async def home(request: Request):
return templates.TemplateResponse("chat.html", {"request": request})
@app.post("/chat")
async def chat(message: ChatMessage):
# Prepare messages
messages = [
{"role": "user", "content": message.message}
]
# Generate response using pipeline
response = pipe(messages)
print(response)
# Extract the response text
if isinstance(response, list):
response_text = response[0].get('generated_text', '')
print(response_text)
else:
response_text = response.get('generated_text', '')
print(response_text)
return {"response": response_text[-1]['content']}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)