Ask-ANRG / main.py
Last commit not found
raw
history blame
1.32 kB
from fastapi import FastAPI, Depends, HTTPException, Query
from transformers import AutoModelForCausalLM, AutoTokenizer
from typing import List
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
app = FastAPI()
app.mount("/static", StaticFiles(directory="static"), name="static")
@app.get("/", response_class=HTMLResponse)
async def read_root():
with open("static/index.html", "r") as f:
content = f.read()
return HTMLResponse(content=content)
# Initialize model and tokenizer
# tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen-14B-Chat-int4")
# model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-14B-Chat-int4").eval()
@app.post("/chat/")
def chat(user_input: str, api_key: str):
# Here you can validate the API key, e.g., check if it exists in your database
# If the API key is not valid, raise an HTTPException
# if not validate_api_key(api_key):
# raise HTTPException(status_code=400, detail="Invalid API key")
# Tokenize the user input and get model's response
# input_ids = tokenizer.encode(user_input, return_tensors="pt")
# output = model.generate(input_ids)
# response = tokenizer.decode(output[0], skip_special_tokens=True)
return {"response": f"user input: {user_input}, api_key: {api_key}"}