Update app.py

app.py

@@ -1,37 +1,20 @@
-import requests
+import gradio as gr
 import json
-import streamlit as st
-
-url = "https://run.cerebrium.ai/pygmalion-6b-webhook/predict"
+import requests
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-payload = json.dumps({"private": "ab48d14619365f6032d4"})
+model_name = 'Pyg'
+tokenizer = AutoTokenizer.from_pretrained("TheBloke/Pygmalion-7B-SuperHOT-8K-GPTQ")
+model = AutoModelForCausalLM.from_pretrained("TheBloke/Pygmalion-7B-SuperHOT-8K-GPTQ")
 
-headers = {
-  'Authorization': 'eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJwcm9qZWN0SWQiOiJwLTRjYmJhNTc0IiwiaWF0IjoxNjkzNzc2Mzg2LCJleHAiOjIwMDkzNTIzODZ9.Puo6VqpcL3iUPNirXGCb31jn42TQ1zho0eaQBHhNhmbv2d5DfqiO2B6U3VlcEnD7JpwUCkMgqfZCQjTkV0a5pTusx15DfXekQBCH_VtDlkm0BsjHvxSICN7RwRQD84xqoJEPe2EuTdfDkQm0Bz18ERSHD8jkUvYGTNDb1FyYXK_yG3qMZuLw8Cpl4l7ivyhLqyxXS_0AvUPHgAeDVISMsKVt4z2nTdcHTGiHNY2rt3INhjLCgnLYY0KGeyBcyhZrMXJyZ3cJfZsGwrsCq6KpEJO_uJAANVhABKKpVzOj163GG8bby19CaSpboezOGZIIvx3G6vRUCUSz4LTRLslNpg',
-  'Content-Type': 'application/json'
-}
+def generate_text(input_text):
+    input_ids = tokenizer.encode(input_text, return_tensors='pt')
+    outputs = model.generate(input_ids, max_length=150, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id)
+    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return text
 
-# Accept user input
-if prompt := st.chat_input("What is up?"):
-    # Display user message in chat message container
-    with st.chat_message("user"):
-        st.markdown(prompt)
-    # Add user message to chat history
-    st.session_state.messages.append({"role": "user", "content": prompt})
-    
-    # Encode the new user input and add end of sentence token
-    inputs = tokenizer.encode(prompt + tokenizer.eos_token, return_tensors="pt")
-    
-    # Generate a response
-    outputs = model.generate(inputs, max_length=50, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id)
-    
-    # Decode the response
-    response = requests.request("POST", url, headers=headers, data=payload)
-    
-    # Display the response in the chat
-    with st.chat_message("bot"):
-        st.markdown(response)
-    # Add bot message to chat history
-    st.session_state.messages.append({"role": "bot", "content": response})
+iface = gr.Interface(fn=generate_text, 
+                      inputs=gr.inputs.Textbox(lines=5, placeholder='Enter text here...'), 
+                      outputs=gr.outputs.Textbox())
 
-print(response.text)
+iface.launch()