Update app.py

app.py

@@ -1,15 +1,13 @@
 import numpy as np
 import streamlit as st
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
-from dotenv import load_dotenv
-import requests
 
-# Load environment variables
-load_dotenv()
-
-# Hugging Face API URL and token
-HUGGINGFACE_API_URL = ["joermd/llma-speedy"]
-HUGGINGFACE_API_TOKEN = os.environ.get('HUGGINGFACEHUB_API_TOKEN')
+# Load the model and tokenizer
+model_name = "joermd/llma-speedy"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
 # Random dog images for error messages
 random_dog = [
@@ -30,7 +28,6 @@ max_token_value = st.sidebar.slider('Select a max_token value', 1000, 9000, 5000
 st.sidebar.button('Reset Chat', on_click=reset_conversation)
 
 # Set the model and display its name
-model_name = "joermd/llma-speedy"
 st.sidebar.write(f"You're now chatting with **{model_name}**")
 st.sidebar.markdown("*Generated content may be inaccurate or false.*")
 
@@ -52,20 +49,14 @@ if prompt := st.chat_input(f"Hi, I'm {model_name}, ask me a question"):
     # Display assistant response
     with st.chat_message("assistant"):
         try:
-            headers = {"Authorization": f"Bearer {HUGGINGFACE_API_TOKEN}"}
-            payload = {
-                "inputs": prompt,
-                "parameters": {"temperature": temp_values, "max_new_tokens": max_token_value}
-            }
-            response = requests.post(HUGGINGFACE_API_URL, headers=headers, json=payload)
-
-            if response.status_code == 200:
-                result = response.json()
-                assistant_response = result.get("generated_text", "No response generated.")
-            else:
-                assistant_response = "Error: Unable to reach the model."
-                st.write(f"Status Code: {response.status_code}")
-                
+            inputs = tokenizer(prompt, return_tensors="pt")
+            outputs = model.generate(
+                inputs.input_ids,
+                max_new_tokens=max_token_value,
+                temperature=temp_values,
+                do_sample=True
+            )
+            assistant_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         except Exception as e:
             assistant_response = "😵‍💫 Connection issue! Try again later. Here's a 🐶:"
             st.image(f'https://random.dog/{random_dog[np.random.randint(len(random_dog))]}')