Chavindu/devops

1from fastapi import FastAPI, HTTPException2from pydantic import BaseModel3from transformers import AutoModelForCausalLM, AutoTokenizer4import os5 6# Set a custom cache directory for Hugging Face7os.environ["TRANSFORMERS_CACHE"] = "/app/cache"8 9app = FastAPI()10 11# Load model and tokenizer12model_name = "Chavindu/rag-model-chav"  # Replace with your model13tokenizer = AutoTokenizer.from_pretrained(model_name)14model = AutoModelForCausalLM.from_pretrained(model_name)15 16# Set the pad_token_id to eos_token_id if not already set17model.config.pad_token_id = model.config.eos_token_id18 19# Define request and response model for FastAPI20class GenerateRequest(BaseModel):21    query: str22 23@app.post("/generate")24async def generate_text(request: GenerateRequest):25    prompt = request.query26    if not prompt:27        raise HTTPException(status_code=400, detail="Prompt is required")28 29    # Tokenize the prompt and generate response30    inputs = tokenizer(prompt, return_tensors="pt")31    output = model.generate(32        inputs["input_ids"],33        max_new_tokens=300,34        temperature=0.6,35        top_p=0.9,36        do_sample=True,37        no_repeat_ngram_size=2,38        eos_token_id=tokenizer.eos_token_id39    )40 41    # Decode and prepare the output42    response = tokenizer.decode(output[0], skip_special_tokens=True)43    generated_text = response.replace(prompt, "").strip()44 45    return {"input_prompt": prompt, "generated_summary": generated_text}46 47if __name__ == "__main__":48    import uvicorn49    uvicorn.run(app, host="0.0.0.0", port=8000)50