Update app.py

app.py

@@ -1,107 +1,96 @@
-import gradio as gr
-import requests
-import json
-import os
-
-from spaces_info import description, examples, initial_prompt_value
-
-API_URL = os.getenv("API_URL")
-HF_API_TOKEN = os.getenv("HF_API_TOKEN")
-
-
-def query(payload):
-    print(payload)
-    response = requests.request("POST", API_URL, json=payload, headers={"Authorization": f"Bearer {HF_API_TOKEN}"})
-    print(response)
-    return json.loads(response.content.decode("utf-8"))
-
-
-def inference(input_sentence, max_length, sample_or_greedy, seed=42):
-    if sample_or_greedy == "Sample":
-        parameters = {
-            "max_new_tokens": max_length,
-            "top_p": 0.9,
-            "do_sample": True,
-            "seed": seed,
-            "early_stopping": False,
-            "length_penalty": 0.0,
-            "eos_token_id": None,
-        }
-    else:
-        parameters = {
-            "max_new_tokens": max_length,
-            "do_sample": False,
-            "seed": seed,
-            "early_stopping": False,
-            "length_penalty": 0.0,
-            "eos_token_id": None,
-        }
+#!/usr/bin/env python
+# coding: utf-8
 
-    payload = {"inputs": input_sentence, "parameters": parameters,"options" : {"use_cache": False} }
+# In[ ]:
 
-    data = query(payload)
 
-    if "error" in data:
-        return (None, None, f"<span style='color:red'>ERROR: {data['error']} </span>")
-
-    generation = data[0]["generated_text"].split(input_sentence, 1)[1]
-    return (
-        before_prompt
-        + input_sentence
-        + prompt_to_generation
-        + generation
-        + after_generation,
-        data[0]["generated_text"],
-        "",
-    )
-
-
-if __name__ == "__main__":
-    demo = gr.Blocks()
-    with demo:
-        with gr.Row():
-            gr.Markdown(value=description)
-        with gr.Row():
-            with gr.Column():
-                text = gr.Textbox(
-                    label="Input",
-                    value=" ",  # should be set to " " when plugged into a real API
-                )
-                tokens = gr.Slider(1, 64, value=32, step=1, label="Tokens to generate")
-                sampling = gr.Radio(
-                    ["Sample", "Greedy"], label="Sample or greedy", value="Sample"
-                )
-                sampling2 = gr.Radio(
-                    ["Sample 1", "Sample 2", "Sample 3", "Sample 4", "Sample 5"],
-                    value="Sample 1",
-                    label="Sample other generations (only work in 'Sample' mode)",
-                    type="index",
-                )
-
-                with gr.Row():
-                    submit = gr.Button("Submit")
-                    load_image = gr.Button("Generate Image")
-            with gr.Column():
-                text_error = gr.Markdown(label="Log information")
-                text_out = gr.Textbox(label="Output")
-                display_out = gr.HTML(label="Image")
-                display_out.set_event_trigger(
-                    "load",
-                    fn=None,
-                    inputs=None,
-                    outputs=None,
-                    no_target=True,
-                    js=js_load_script,
-                )
-        with gr.Row():
-            gr.Examples(examples=examples, inputs=[text, tokens, sampling, sampling2])
-
-        submit.click(
-            inference,
-            inputs=[text, tokens, sampling, sampling2],
-            outputs=[display_out, text_out, text_error],
-        )
-
-        load_image.click(fn=None, inputs=None, outputs=None, _js=js_save)
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import gradio as gr
+import re
+
+def cleaning_history_tuple(history):
+    s=sum(history,())
+    s=list(s)
+    s2=""
+    for i in s:
+        i=re.sub("\n", '', i)
+        i=re.sub("<p>", '', i)
+        i=re.sub("</p>", '', i)
+        s2=s2+i+'\n'
+    return s2
+
+def ai_output(string1,string2):
+    a1=len(string1)
+    a2=len(string2)
+    string3=string2[a1:]
+    sub1="A:"
+    sub2="User"
+    #sub3="\n"
+    try:
+        try:
+            idx1=string3.index(sub1)
+            response=string3[:idx1]
+            return response
+            
+        except:
+            idx1=string3.index(sub2)
+            response=string3[:idx1]
+            return response
+    except:
+        return string3
+
+model4 = AutoModelForCausalLM.from_pretrained("bigscience/bloom-3b")
+tokenizer4 = AutoTokenizer.from_pretrained("bigscience/bloom-3b")
+
+def predict(input,initial_prompt, temperature=0.7,top_p=1,top_k=5,max_tokens=64,no_repeat_ngram_size=1,num_beams=6,do_sample=True, history=[]):
+	
+    s = cleaning_history_tuple(history)
+    
+    s = s+ "\n"+ "User: "+ input + "\n" + "Assistant: "
+    s2=initial_prompt+" " + s
+    
+    input_ids = tokenizer4.encode(str(s2), return_tensors="pt")
+    response = model4.generate(input_ids, min_length = 10,
+                         max_new_tokens=int(max_tokens),
+                         top_k=int(top_k),
+                         top_p=float(top_p),
+                         temperature=float(temperature),
+                         no_repeat_ngram_size=int(no_repeat_ngram_size),
+                         num_beams = int(num_beams),
+                         do_sample = bool(do_sample),
+                         )
+    
+  
+    response2 = tokenizer4.decode(response[0])
+    print("Response after decoding tokenizer: ",response2)
+    print("\n\n")
+    response3=ai_output(s2,response2)
+    
+    input="User: "+input
+    response3="Assistant: "+ response3
+    history.append((input, response3))
+
+    return history, history
+
+#gr.Interface(fn=predict,title="BLOOM-3b",
+#             inputs=["text","text","text","text","text","text","text","text","text",'state'],
+#            
+#             outputs=["chatbot",'state']).launch()
+             
+
+gr.Interface(inputs=[gr.Textbox(label="input", lines=1, value=""),
+                     gr.Textbox(label="initial_prompt", lines=1, value=prompt),
+                     gr.Textbox(label="temperature", lines=1, value=0.7),
+                     gr.Textbox(label="top_p", lines=1, value=1),
+                     gr.Textbox(label="top_k", lines=1, value=5),
+                     gr.Textbox(label="max_tokens", lines=1, value=64),
+                     gr.Textbox(label="no_repeat_ngram_size", lines=1, value=1),
+                     gr.Textbox(label="num_beams", lines=1, value=6),
+                     gr.Textbox(label="do_sample", lines=1, value="True"), 'state'],
+             fn=predict, title="OPT-6.7B", outputs=["chatbot",'state']
+
+             #inputs=["text","text","text","text","text","text","text","text","text",'state'],
+
+             ).launch()            
 
-    demo.launch()