Update app.py

app.py

@@ -1,14 +1,28 @@
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import gradio as gr
+from fastapi import FastAPI
+from pydantic import BaseModel
+import uvicorn
 
+# Load model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/deepseek-llm-1.8b-chat")
 model = AutoModelForCausalLM.from_pretrained("deepseek-ai/deepseek-llm-1.8b-chat")
 
-def chat(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
+# Initialize FastAPI app
+app = FastAPI()
+
+# Define input schema
+class RequestBody(BaseModel):
+    prompt: str
+
+# Define the model inference function
+@app.post("/predict")
+async def predict(request: RequestBody):
+    inputs = tokenizer(request.prompt, return_tensors="pt")
     outputs = model.generate(**inputs, max_new_tokens=100)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return {"response": result}
 
-demo = gr.Interface(fn=chat, inputs="text", outputs="text")
-demo.launch()
+# For testing locally (not needed for Hugging Face)
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)