Add FastAPI for handling POST requests

app.py

@@ -1,29 +1,31 @@
 import streamlit as st
 from transformers import pipeline
-from fastapi import FastAPI
-from pydantic import BaseModel
-import uvicorn
+import subprocess
+import sys
 
-app = FastAPI()
 
+# Ensure PyTorch is installed
+def install_pytorch():
+    subprocess.check_call([sys.executable, "-m", "pip", "install", "torch"])
 
-class Input(BaseModel):
-    input: str
 
+try:
+    import torch
+except ImportError:
+    install_pytorch()
 
-@st.cache(allow_output_mutation=True)
+st.title("Hugging Face Model Demo")
+
+
+@st.cache_resource
 def load_model():
     return pipeline("text-generation", model="klyang/MentaLLaMA-chat-7B")
 
 
 model = load_model()
 
-
-@app.post("/")
-def generate_text(input: Input):
-    result = model(input.input)
-    return {"generated_text": result[0]["generated_text"]}
-
-
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+user_input = st.text_input("Enter your text:")
+if user_input:
+    with st.spinner("Generating response..."):
+        result = model(user_input)
+        st.success(result[0]["generated_text"])