Update app.py

app.py

@@ -1,64 +1,47 @@
-from huggingface_hub import InferenceClient
+import torch
+from transformers import GPT2Tokenizer, GPT2LMHeadModel
 import gradio as gr
-import json
 
-client = InferenceClient(
-    "mistralai/Mistral-7B-Instruct-v0.3"
+# Check if a GPU is available and use it, otherwise use CPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load the pre-trained model and tokenizer from the saved directory
+model_path = "blexus_pretrained_test"
+tokenizer = GPT2Tokenizer.from_pretrained(model_path)
+model = GPT2LMHeadModel.from_pretrained(model_path).to(device)
+
+# Set model to evaluation mode
+model.eval()
+
+# Function to generate text based on input prompt
+def generate_text(prompt):
+    # Tokenize and encode the input prompt
+    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
+
+    # Generate continuation
+    with torch.no_grad():
+        generated_ids = model.generate(
+            input_ids, 
+            max_length=50,                # Maximum length of generated text
+            num_return_sequences=1,        # Generate 1 sequence
+            pad_token_id=tokenizer.eos_token_id,  # Use EOS token for padding
+            do_sample=True,                # Enable sampling
+            top_k=50,                      # Top-k sampling
+            top_p=0.95                     # Nucleus sampling
+        )
+
+    # Decode the generated text
+    generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+    return generated_text
+
+# Create a Gradio interface
+interface = gr.Interface(
+    fn=generate_text,                      # Function to call when interacting with the UI
+    inputs="text",                         # Input type: Single-line text
+    outputs="text",                        # Output type: Text (the generated output)
+    title="Quble Text Generation",         # Title of the UI
+    description="Enter a prompt to generate text using Quble."  # Simple description
 )
 
-DATABASE_PATH = "database.json"
-
-def load_database():
-    try:
-        with open(DATABASE_PATH, "r") as file:
-            return json.load(file)
-    except FileNotFoundError:
-        return {}
-
-def save_database(database):
-    with open(DATABASE_PATH, "w") as file:
-        json.dump(database, file)
-
-def format_prompt(message, history):
-    prompt = "<s>"
-    for user_prompt, bot_response in history:
-        prompt += f"[INST] {user_prompt} [/INST]"
-        prompt += f" {bot_response}</s> "
-    prompt += f"[INST] {message} [/INST]"
-    return prompt
-
-def generate(
-    prompt, history, temperature=0.9, max_new_tokens=4096, top_p=0.9, repetition_penalty=1.2,
-):
-    database = load_database()  # Load the database
-    temperature = float(temperature)
-    if temperature < 1e-2:
-        temperature = 1e-2
-    top_p = float(top_p)
-
-    formatted_prompt = format_prompt(prompt, history)
-    if formatted_prompt in database:
-        response = database[formatted_prompt]
-    else:
-        response = client.text_generation(formatted_prompt, details=True, return_full_text=False)
-        response_text = response.generated_text
-        database[formatted_prompt] = response_text
-        save_database(database)  # Save the updated database
-
-    yield response_text
-
-css = """
-  #mkd {
-    height: 500px; 
-    overflow: auto; 
-    border: 1px solid #ccc; 
-  }
-"""
-
-with gr.Blocks(css=css) as demo:
-    gr.ChatInterface(
-        generate,
-        examples=[["What is the secret to life?"], ["Write me a recipe for pancakes."], ["Write a short story about Paris."]]
-    )
-
-demo.launch(debug=True)
+# Launch the Gradio app
+interface.launch()