test interface

app.py

@@ -1,30 +1,25 @@
 import gradio as gr
-import os
-import json
-import openai
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from run_llm import model_mapping, fastchat  # Import the necessary function from run_llm.py
 
-# Set your OpenAI API key
-openai.api_key = "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"
+# Load pre-trained GPT-3.5 model and tokenizer (you can replace this with your model)
+model_name = "EleutherAI/gpt-neo-2.7B"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
-def generate_text(input_text, model, prompt_type):
-    # Use the fastchat function from run_llm.py
-    outputs = fastchat(input_text, model, prompt_type)
-    return outputs
+def generate_text(input_text, max_length=50):
+    input_ids = tokenizer.encode(input_text, return_tensors="pt")
+    output = model.generate(input_ids, max_length=max_length, num_return_sequences=1)
+    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+    return generated_text
 
+# Create a Gradio interface
 iface = gr.Interface(
-    fn=generate_text,
-    inputs=[
-        gr.Textbox("input_text", label="Input Text"),
-        gr.Dropdown(
-            list(model_mapping.keys()),
-            label="Model"
-        ),
-        gr.Radio([1, 2], label="Prompt Type"),
-    ],
-    outputs=gr.Textbox("output_text", label="Generated Text")
+    fn=generate_text,  # Your text generation function
+    inputs=gr.Textbox(text="Enter text here..."),  # Text input field
+    outputs=gr.Textbox(),  # Display generated text
+    live=True  # Real-time updates
 )
 
+# Launch the interface
 iface.launch()