app.py

from models import MODEL_IDS
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from transformers import AutoModelForCausalLM, AutoTokenizer
from langchain.memory import ConversationBufferMemory
import torch
import asyncio

# Initialize FastAPI app
app = FastAPI()

# Load Hugging Face model and tokenizer
MODEL_NAME = MODEL_IDS["llama_3.2_1b_instruct"]
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

# Initialize LangChain memory for conversation history
memory = ConversationBufferMemory()

# Pydantic model for request body


class PromptRequest(BaseModel):
    prompt: str

# Function to generate streaming responses


async def generate_response(prompt: str):
    # Add the prompt to memory
    memory.chat_memory.add_user_message(prompt)

    # Tokenize the input prompt
    inputs = tokenizer(prompt, return_tensors="pt",
                       max_length=512, truncation=True)

    # Generate response tokens incrementally
    with torch.no_grad():
        for _ in range(50):  # Adjust max tokens as needed
            outputs = model.generate(
                inputs["input_ids"],
                max_length=inputs["input_ids"].shape[1] + 1,
                pad_token_id=tokenizer.eos_token_id,
                do_sample=True,
                top_k=50,
            )
            new_token = outputs[0, -1].item()
            if new_token == tokenizer.eos_token_id:
                break

            # Decode the new token and yield it
            decoded_token = tokenizer.decode(
                [new_token], skip_special_tokens=True)
            yield decoded_token

            # Update inputs for the next iteration
            inputs["input_ids"] = torch.cat(
                [inputs["input_ids"], torch.tensor([[new_token]])], dim=1)

    # Add the model's response to memory
    full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    memory.chat_memory.add_ai_message(full_response)

# FastAPI endpoint for streaming responses


@app.post("/stream-chat/")
async def stream_chat(request: PromptRequest):
    if not request.prompt:
        raise HTTPException(status_code=400, detail="Prompt cannot be empty")

    return StreamingResponse(generate_response(request.prompt), media_type="text/plain")

# FastAPI endpoint to get conversation history


@app.get("/conversation-history/")
async def get_conversation_history():
    return {"history": memory.load_memory_variables({})}

# Run the FastAPI app
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)