rr19tech/llama_index

1from fastapi import FastAPI2import onnxruntime_genai as og3import numpy as np4import json5 6app = FastAPI()7 8# Global model loading to prevent reload on every request9model = og.Model("./model/cpu_and_mobile/cpu-int4-rtn-block-32")10#testing some thread options to make it faster11import os12# Force only 2 threads to match HF Free Tier CPUs13os.environ["OMP_NUM_THREADS"] = "2"14os.environ["MKL_NUM_THREADS"] = "2"15 16tokenizer = og.Tokenizer(model)17 18@app.get("/chat")19async def chat(user_input: str, retrieved_context: str):20    # Tool definitions must be injected into the system prompt for 0.5B models21    22    system_prompt="""You are a concise, knowledgeable assistant. Respond to all questions in 2-3 well-formed, accurate sentences. 23    Avoid filler phrases ("Certainly!", "Great question!"), conversational apologies, and overly detailed descriptions. 24    Prioritize directness and clarity. Use the provided context if applicable."""25    26    messages = [27        {"role": "system", "content": f"{system_prompt}"},28        {"role": "user", "content": f"Context: {retrieved_context}\n\nQuestion: {user_input}"}29    ]30 31    # 3. Use apply_chat_template to get the correctly formatted string32    # This automatically adds markers like <|user|>, [INST], etc.33    prompt = tokenizer.apply_chat_template(json.dumps(messages), add_generation_prompt=True)34             35    params = og.GeneratorParams(model)36    tokens = tokenizer.encode(prompt)37 38    #params = og.GeneratorParams(model, tokens)39    #params.set_search_options(max_length=512, temperature=0.1) # Low temp for better tool accuracy40    params.set_search_options(max_length=5000, temperature=0.8, top_p=0.9, do_sample=True, repetition_penalty=1.1)41    generator = og.Generator(model, params)42    generator.append_tokens(np.array(tokens, dtype=np.int32))43    output = ""44    45    # 1. Run the generation loop46    while not generator.is_done():47        generator.generate_next_token()48 49    # 2. Get the full sequence of token IDs50    output_tokens = generator.get_sequence(0)51    input_length = len(tokens) # saving the input tokens length, so we can remove all the extra text before the new token52    new_tokens_only = output_tokens[input_length:] # getting the new tokens only and we will decode only this53 54    # 3. Decode back to text and print55    #response = tokenizer.decode(output_tokens)56    response = tokenizer.decode(new_tokens_only)57    print(response)58 59    return response60 61#adding a new method for structured output62from pydantic import BaseModel, Field63from typing import List64from enum import Enum65 66#identifying potential db operations67class DbIntent(str, Enum):68    CREATE = "create"69    READ = "read"70    UPDATE = "update"71    DELETE = "delete"72 73# 1. Define your Pydantic schema74class JobDescription(BaseModel):75    job_description: str = Field(description="The full job description")76    title: str77    keywords: List[str]78    summary: str = Field(description="A brief 1-sentence summary")79    company: str80    intent: DbIntent81 82schema = JobDescription.model_json_schema()83print(schema)84import json85schema_str = json.dumps(schema)86 87@app.get("/structured_chat")88async def chat(user_input: str, retrieved_context: str):89    # Tool definitions must be injected into the system prompt for 0.5B models90    91    system_prompt="You are a helpful assistant with a professional tone."92    93    messages = [94        {"role": "system", "content": f"{system_prompt}"},95        {"role": "user", "content": f"Question: {user_input}"}96    ]97    98    # 3. Use apply_chat_template to get the correctly formatted string99    # This automatically adds markers like <|user|>, [INST], etc.100    prompt = tokenizer.apply_chat_template(json.dumps(messages), add_generation_prompt=True)101             102    params = og.GeneratorParams(model)103    tokens = tokenizer.encode(prompt)104    105    #params = og.GeneratorParams(model, tokens)106    #params.set_search_options(max_length=512, temperature=0.1) # Low temp for better tool accuracy107    params.set_search_options(max_length=5000, temperature=0.1)108    params.set_guidance("json_schema", schema_str) #below does not work109    #params.set_logits_processor(110        #og.LogitsProcessor.from_json_schema(model, schema)111    #)112    generator = og.Generator(model, params)113    generator.append_tokens(np.array(tokens, dtype=np.int32))114    output = ""115    116    # 1. Run the generation loop117    while not generator.is_done():118        generator.generate_next_token()119    120    # 2. Get the full sequence of token IDs121    output_tokens = generator.get_sequence(0)122    input_length = len(tokens) # saving the input tokens length, so we can remove all the extra text before the new token123    new_tokens_only = output_tokens[input_length:] # getting the new tokens only and we will decode only this124    125    # 3. Decode back to text and print126    #response = tokenizer.decode(output_tokens)127    response = tokenizer.decode(new_tokens_only)128    print(response)129    130    return response131