Update app.py

app.py

@@ -1,31 +1,32 @@
-import gradio as gr
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer  # Import AutoTokenizer
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+import uvicorn
+
+app = FastAPI()
 
 try:
-    # Load the model
     model = AutoModelForCausalLM.from_pretrained("petertill/cordia-a6")
-
-    # Load the tokenizer
     tokenizer = AutoTokenizer.from_pretrained("petertill/cordia-a6")
-
-    # Create the pipeline with both model and tokenizer
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer) 
-
+    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
     print("Model and tokenizer loaded successfully!")
 
-    def generate_text(prompt):
-        output = pipe(prompt, max_length=200)[0]['generated_text']
-        return output
+    class GenerateRequest(BaseModel):
+        prompt: str
 
-    iface = gr.Interface(
-        fn=generate_text,
-        inputs="text",
-        outputs="text",
-        title="Cordia-A6 Text Generation",
-        description="Generate text using the Cordia-A6 model."
-    )
+    class GenerateResponse(BaseModel):
+        generated_text: str
 
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+    @app.post("/generate", response_model=GenerateResponse)
+    async def generate(request: GenerateRequest):
+        try:
+            output = pipe(request.prompt, max_length=200)[0]['generated_text']
+            return GenerateResponse(generated_text=output)
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e))
 
 except Exception as e:
-    print(f"Error: {e}")
+    print(f"Error: {e}")
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)