zerovic/phi-3-mini-4k-instruct

1import torch2from transformers import AutoTokenizer, AutoModelForCausalLM3from fastapi import FastAPI4from pydantic import BaseModel5 6app = FastAPI()7 8# ✅ Phi-3 model9MODEL_NAME = "microsoft/Phi-3-mini-4k-instruct"10 11tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)12 13model = AutoModelForCausalLM.from_pretrained(14    MODEL_NAME,15    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float3216)17 18class RequestData(BaseModel):19    inputs: str20 21 22def generate_text(prompt):23    # ✅ Add a System Message to enforce "Human-like" brevity24    messages = [25        {26            "role": "system", 27            "content": "You are a concise assistant. Answer the user's question directly. If there is a typo in the question, correct it silently and provide the answer. Do not give unsolicited details."28        },29        {"role": "user", "content": prompt}30    ]31 32    formatted_prompt = tokenizer.apply_chat_template(33        messages,34        tokenize=False,35        add_generation_prompt=True36    )37 38    inputs = tokenizer(formatted_prompt, return_tensors="pt")39    40    # Store the length of the input tokens41    input_length = inputs.input_ids.shape[1]42 43    with torch.no_grad():44        output = model.generate(45            **inputs,46            max_new_tokens=200,47            do_sample=True,48            temperature=0.7,49            top_p=0.9,50            repetition_penalty=1.1,51            pad_token_id=tokenizer.eos_token_id52        )53 54    # ✅ FIX: Slice the output to exclude the input tokens55    # output[0] is the full sequence; [input_length:] takes everything AFTER the prompt56    new_tokens = output[0][input_length:]57    58    result = tokenizer.decode(new_tokens, skip_special_tokens=True)59 60    return result.strip()61 62 63@app.post("/generate")64async def generate(request: RequestData):65 66    text = generate_text(request.inputs)67 68    return {69        "data": [text]70    }