krupakar-reddy's picture
Update main.py
03affe6 verified
from fastapi import FastAPI, Request
from pydantic import BaseModel
from llama_cpp import Llama
model_path = "meta-llama-3-8b-instruct.Q4_K_M.gguf"
llm = Llama(
model_path = model_path,
n_ctx = 4096,
n_threads = 4,
)
app = FastAPI()
class DSAPrompt(BaseModel):
prompt: str
@app.post("/solve_dsa_problem")
async def solve_dsa_problem(item: DSAPrompt):
system_prompt = '''
You are a Data Structures and Algorithm problem solver. You are given the following problem and you need to solve it.
Give a detailed explanation of the solution approach and the code in C++, Java, or Python.
If the input is not a DSA problem, politely refuse their request and reinsist to provide a DSA problem.
'''
res = llm.create_chat_completion(
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": item.prompt}
],
temperature = 0.7,
max_tokens = 2048,
)
return res