Update app.py

app.py

@@ -1,33 +1,64 @@
-from fastapi import FastAPI, Request
-from transformers import pipeline
-import json
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from huggingface_hub import InferenceClient
 import os
+import gradio as gr
 
 app = FastAPI()
 
-# Set your Hugging Face API token
-huggingface_token = os.environ.get("huggingface_token")
+client = InferenceClient(model="HuggingFaceH4/zephyr-7b-beta", token=os.getenv("huggingface_token"))
 
-# Load the model
-generator = pipeline("text-generation", model="EleutherAI/gpt-neo-2.7B")
+class Message(BaseModel):
+    input: str
+    history: list
 
-conversation_history = []
+@app.get("/")
+async def home():
+    return {"message": "Welcome to the chatbot API!"}
 
 @app.post("/chat")
-async def chat(message: dict):
-    global conversation_history
-    # Extract user message
-    user_message = message["message"]
-    conversation_history.append({"role": "user", "content": user_message})
+async def chat(message: Message):
+    try:
+        input_message = message.input
+        history = message.history
+        
+        # Call the chatbot function
+        response = chatbot(input_message, history)
+        
+        return {"response": response}
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
-    messages = [msg["content"] for msg in conversation_history]
+def chatbot(input, history):
+    try:
+        # Call your chatbot function here
+        message = [{"role": "user", "content": input}]
+        history = [{"role": "system", "content": "You are a helpful assistant."}]
+        messages = history + message
+        
+        output = client.chat_completion(
+            messages=messages,
+            max_tokens=256,
+            temperature=0.7
+        )
+        
+        history = history + [{"role": "assistant", "content": output.choices[0].message.content}]
+        
+        return output.choices[0].message.content
+    
+    except Exception as e:
+        raise Exception(str(e))
 
-    # Add a system message to instruct the model
-    messages.insert(0, "You are a story writer. Whatever the prompt, you always write a short story of 30 words.")
+# Define the Gradio chat interface
+def gradio_chat(input, history):
+    response = chatbot(input, history)
+    return response
 
-    # Generate response using Hugging Face model
-    reply = generator(messages, max_length=30, do_sample=False)[0]["generated_text"]
+# Define Gradio inputs and outputs
+inputs = [gr.Textbox(lines=5, label="Input"), gr.Textbox(lines=5, label="History")]
+output = gr.Textbox(label="Response")
 
-    conversation_history.append({"role": "assistant", "content": reply})
+# Create Gradio interface
+gr.Interface(fn=gradio_chat, inputs=inputs, outputs=output, title="Chatbot").launch()
 
-    return reply