Rename inference.py to app.py

inference.py → app.py

@@ -1,6 +1,7 @@
 import os
 from dataclasses import dataclass, asdict
 from ctransformers import AutoModelForCausalLM, AutoConfig
+import gradio as gr
 
 
 @dataclass
@@ -39,6 +40,14 @@ def generate(
     )
 
 
+def generate_response(user_input):
+    generator = generate(llm, generation_config, user_input.strip())
+    response = ""
+    for word in generator:
+        response += word
+    return response
+
+
 if __name__ == "__main__":
     config = AutoConfig.from_pretrained(
         "teknium/Replit-v2-CodeInstruct-3B", context_length=2048
@@ -59,16 +68,17 @@ if __name__ == "__main__":
         reset=True,  # reset history (cache)
         stream=True,  # streaming per word/token
         threads=int(os.cpu_count() / 6),  # adjust for your CPU
-        stop=["<|endoftext|>"],
+        stop=[""],
     )
 
     user_prefix = "[user]: "
-    assistant_prefix = f"[assistant]:"
+    assistant_prefix = f"[assistant]: "
 
-    while True:
-        user_prompt = input(user_prefix)
-        generator = generate(llm, generation_config, user_prompt.strip())
-        print(assistant_prefix, end=" ", flush=True)
-        for word in generator:
-            print(word, end="", flush=True)
-        print("")
+    iface = gr.Interface(
+        fn=generate_response,
+        inputs=gr.inputs.Textbox(label=user_prefix),
+        outputs=gr.outputs.Textbox(label=assistant_prefix),
+        title="Chat with Assistant",
+        description="Ask any question and get a response from the Assistant!",
+    )
+    iface.launch()