Update app.py

app.py

@@ -1,64 +1,153 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import os
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("Ansah-AI/E1-4BIT-GGUF")
+# Check if CUDA is available
+if torch.cuda.is_available():
+    print(f"Using GPU: {torch.cuda.get_device_name(0)}")
+    device = "cuda"
+else:
+    print("GPU not available, using CPU")
+    device = "cpu"
 
+# Model constants
+MODEL_ID = "Ansah-AI/E1-4BIT-GGUF"
+MODEL_REVISION = "main"  # Change this if you need a specific revision
+    
+# Function to download and load the model and tokenizer
+def load_model():
+    print(f"Loading model: {MODEL_ID}")
+    
+    # Load model with 4-bit quantization
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        revision=MODEL_REVISION,
+        device_map="auto",
+        load_in_4bit=True,  # Enable 4-bit quantization
+        trust_remote_code=True
+    )
+    
+    # Load tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+    
+    print("Model and tokenizer loaded successfully!")
+    
+    return model, tokenizer
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
+# Function to generate text from the model
+def generate_text(prompt, max_length=256, temperature=0.7, top_p=0.9, top_k=40):
+    # Ensure the model and tokenizer are loaded
+    global model, tokenizer
+    
+    # Print generation parameters for debugging
+    print(f"Generating with parameters: max_length={max_length}, temp={temperature}, top_p={top_p}, top_k={top_k}")
+    
+    # Create the text generation pipeline
+    text_generator = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        device_map="auto"
+    )
+    
+    # Generate the text
+    generation_config = {
+        "max_length": max_length,
+        "temperature": temperature,
+        "top_p": top_p,
+        "top_k": top_k,
+        "num_return_sequences": 1,
+        "do_sample": temperature > 0.1,  # Use sampling if temperature is significant
+        "pad_token_id": tokenizer.eos_token_id
+    }
+    
+    try:
+        result = text_generator(
+            prompt,
+            **generation_config
+        )
+        
+        # Return the generated text
+        return result[0]["generated_text"]
+    
+    except Exception as e:
+        return f"Error generating text: {str(e)}"
 
+# Main function to create and run the Gradio interface
+def main():
+    # Load the model and tokenizer
+    global model, tokenizer
+    model, tokenizer = load_model()
+    
+    # Create the Gradio interface
+    with gr.Blocks(title="E1-4BIT-GGUF Model Interface") as demo:
+        gr.Markdown("# E1-4BIT-GGUF Model Interface")
+        gr.Markdown("Enter your prompt below to generate text using the Ansah-AI/E1-4BIT-GGUF model.")
+        
+        with gr.Row():
+            with gr.Column(scale=4):
+                prompt_input = gr.Textbox(
+                    label="Prompt",
+                    placeholder="Enter your prompt here...",
+                    lines=5
+                )
+            
+            with gr.Column(scale=1):
+                max_length = gr.Slider(
+                    minimum=64,
+                    maximum=2048,
+                    value=256,
+                    step=32,
+                    label="Max Length"
+                )
+                
+                temperature = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.5,
+                    value=0.7,
+                    step=0.1,
+                    label="Temperature"
+                )
+                
+                top_p = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    value=0.9,
+                    step=0.05,
+                    label="Top P"
+                )
+                
+                top_k = gr.Slider(
+                    minimum=1,
+                    maximum=100,
+                    value=40,
+                    step=1,
+                    label="Top K"
+                )
+        
+        generate_button = gr.Button("Generate")
+        output_text = gr.Textbox(label="Generated Text", lines=10)
+        
+        # Set up the button click event
+        generate_button.click(
+            fn=generate_text,
+            inputs=[prompt_input, max_length, temperature, top_p, top_k],
+            outputs=output_text
+        )
+        
+        # Add examples
+        gr.Examples(
+            examples=[
+                ["Write a short story about a space explorer discovering a new planet."],
+                ["Explain quantum computing to a high school student."],
+                ["Create a recipe for a chocolate cake."]
+            ],
+            inputs=prompt_input
+        )
+    
+    # Launch the interface
+    demo.launch(share=True)  # share=True creates a public link
 
 if __name__ == "__main__":
-    demo.launch()
+    main()