Create main.py

main.py

@@ -0,0 +1,65 @@
+import torch
+from fastapi import FastAPI
+from fastapi.responses import StreamingResponse
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from threading import Thread
+
+app = FastAPI()
+
+MODEL_ID = "AshokGakr/model-tiny"
+
+print("Loading model...")
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.float32,
+    low_cpu_mem_usage=True
+).to(device)
+
+model.eval()
+
+print("Model loaded on", device)
+
+
+def generate_stream(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+    streamer = TextIteratorStreamer(
+        tokenizer,
+        skip_prompt=True,
+        skip_special_tokens=True
+    )
+
+    generation_kwargs = dict(
+        **inputs,
+        max_new_tokens=120,
+        temperature=0.7,
+        top_p=0.9,
+        repetition_penalty=1.1,
+        do_sample=True,
+        streamer=streamer
+    )
+
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    for new_text in streamer:
+        yield new_text
+
+
+@app.post("/chat")
+async def chat(data: dict):
+    system_prompt = data.get("system", "You are a helpful AI assistant.")
+    history = data.get("history", "")
+    message = data.get("message", "")
+
+    full_prompt = f"{system_prompt}\n{history}\nUser: {message}\nAssistant:"
+
+    return StreamingResponse(
+        generate_stream(full_prompt),
+        media_type="text/plain"
+    )