vis / api.py
Guhanselvam's picture
Create api.py
386fc65 verified
raw
history blame contribute delete
688 Bytes
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from llama_cpp import Llama # Change this import as per your LLaMA library
app = FastAPI()
# Initialize LLaMA model
model = Llama(model_path="path/to/llama-3.2-model.bin") # Specify the correct path to your model
class LlamaRequest(BaseModel):
prompt: str
@app.post("/llama/")
async def get_llama_response(request: LlamaRequest):
try:
response = model.generate(request.prompt) # Call the LLaMA model
return {"response": response}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
# To run the API, use 'uvicorn api:app --reload' in the terminal.