main.py

from fastapi import FastAPI
from pydantic import BaseModel
from huggingface_hub import InferenceClient

app = FastAPI()

# Initialize the inference client for the AI model
client = InferenceClient("nvidia/Llama-3.1-Nemotron-70B-Instruct-HF")

class CourseRequest(BaseModel):
    course_name: str

def format_prompt(course_name: str):
    return f"Please generate a detailed roadmap for the course '{course_name}'. Include key topics, recommended resources, and a learning timeline."

def generate_roadmap(item: CourseRequest):
    prompt = format_prompt(item.course_name)

    # You can adjust these parameters as needed
    generate_kwargs = {
        "temperature": 0.7,
        "max_new_tokens": 150,
        "top_p": 0.9,
        "repetition_penalty": 1.1,
    }

    # Call the model to generate the roadmap
    stream = client.text_generation(prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
    output = ""

    for response in stream:
        output += response.token.text
    
    return output

@app.post("/generate/")
async def generate_roadmap_endpoint(course_request: CourseRequest):
    roadmap = generate_roadmap(course_request)
    return {"roadmap": roadmap}