Update app.py

app.py

@@ -1,24 +1,21 @@
-import os
-from huggingface_hub import InferenceClient
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-# Set your Hugging Face API key as an environment variable for security
-os.environ["HUGGINGFACE_API_KEY"] = "hf_xxxxxxxxxxxxxxxxxxxxxxxx"
+# Load model and tokenizer
+model_name = "inceptionai/jais-13b"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
 
-# Initialize the client without the 'provider' argument
-client = InferenceClient()
+# Define chatbot function
+def chat_with_jais(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    outputs = model.generate(**inputs, max_length=512)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-messages = [
-    {
-        "role": "user",
-        "content": "What is the capital of France?"
-    }
-]
+# Gradio Interface
+interface = gr.Interface(fn=chat_with_jais, inputs="text", outputs="text", title="JAIS-13B Chatbot")
 
-# Use the chat API with the specified model and messages
-completion = client.chat.completions.create(
-    model="meta-llama/Llama-2-7b-chat-hf", 
-    messages=messages, 
-    max_tokens=500
-)
-
-print(completion.choices[0].message)
+# Launch the app
+if __name__ == "__main__":
+    interface.launch()