Update main.py

main.py

@@ -1,22 +1,26 @@
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from transformers import pipeline
+
+# Load a chat model (you can choose a bigger one)
+generator = pipeline('text-generation', model='gpt2')  # gpt2 is small; for better, use GPT-Neo
 
 app = FastAPI()
 
+# Allow CORS for your frontend
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
-    allow_credentials=True,
     allow_methods=["*"],
-    allow_headers=["*"],
+    allow_headers=["*"]
 )
 
-@app.get("/")
-def root():
-    return {"message": "Backend is working"}
+class Prompt(BaseModel):
+    prompt: str
 
 @app.get("/chat")
 def chat(prompt: str):
-    return {
-        "reply": f"You said: {prompt}. This is your AI running on the internet."
-    }
+    # Generate response
+    response = generator(prompt, max_length=100, do_sample=True)[0]['generated_text']
+    return {"reply": response}