Update app.py

app.py

@@ -1,15 +1,7 @@
-from fastapi import FastAPI, Request
-from pydantic import BaseModel
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-import os
+import gradio as gr
 
-app = FastAPI()
-
-class UserInput(BaseModel):
-    prompt: str
-
-# Load model and tokenizer
 model_name = "NousResearch/Nous-Hermes-llama2-13b"
 tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
 model = AutoModelForCausalLM.from_pretrained(
@@ -21,14 +13,16 @@ model = AutoModelForCausalLM.from_pretrained(
 
 generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-@app.post("/chat")
-async def chat(data: UserInput):
+def chat(prompt):
     output = generator(
-        data.prompt,
+        prompt,
         max_new_tokens=512,
         do_sample=True,
         temperature=0.8,
-        top_k=50,
+        top_k=60,
         top_p=0.95
     )
-    return {"response": output[0]["generated_text"]}
+    return output[0]["generated_text"]
+
+demo = gr.Interface(fn=chat, inputs="text", outputs="text")
+demo.launch()