Update app.py

app.py

@@ -1,25 +1,20 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from transformers import pipeline
 
-# Define available models
-MODELS = {
-    "Zephyr 7B Beta": "HuggingFaceH4/zephyr-7b-beta",
-    "GPT-2": "gpt2",
-    "GPT-2 Medium": "gpt2-medium", 
-    "DistilGPT-2": "distilgpt2",
-    "German GPT-2": "german-nlp-group/german-gpt2",
-    "German Wechsel GPT-2": "benjamin/gpt2-wechsel-german",
-    "T5 Base": "t5-base",
-    "T5 Large": "t5-large"
-}
-
-def create_inference_client(model_name):
-    """Create an InferenceClient for the selected model."""
-    try:
-        return InferenceClient(model_name)
-    except Exception as e:
-        print(f"Error creating client for {model_name}: {e}")
-        return None
+# Define all pipelines
+def load_pipelines():
+    pipelines = {
+        "GPT-2 Original": pipeline("text-generation", model="gpt2"),
+        "GPT-2 Medium": pipeline("text-generation", model="gpt2-medium"),
+        "DistilGPT-2": pipeline("text-generation", model="distilgpt2"),
+        "German GPT-2": pipeline("text-generation", model="german-nlp-group/german-gpt2"),
+        "German Wechsel GPT-2": pipeline("text-generation", model="benjamin/gpt2-wechsel-german"),
+        "T5 Base": pipeline("text-generation", model="t5-base"),
+        "T5 Large": pipeline("text-generation", model="t5-large"),
+        "Text Classification": pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english"),
+        "Sentiment Analysis": pipeline("sentiment-analysis", model="nlptown/bert-base-multilingual-uncased-sentiment")
+    }
+    return pipelines
 
 def respond(
     message,
@@ -30,45 +25,65 @@ def respond(
     temperature,
     top_p,
 ):
-    """Generate response using selected model."""
-    # Create client for selected model
-    client = create_inference_client(MODELS[model_name])
+    # Load pipelines
+    pipelines = load_pipelines()
+    pipe = pipelines.get(model_name)
     
-    if not client:
-        return "Error: Could not create model client."
+    if not pipe:
+        return "Error: Model not found."
 
-    # Prepare chat history
-    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})
+    # For text generation models
+    if model_name in ["GPT-2 Original", "GPT-2 Medium", "DistilGPT-2", 
+                      "German GPT-2", "German Wechsel GPT-2", 
+                      "T5 Base", "T5 Large"]:
+        # Prepare full prompt
+        full_history = ' '.join([f"{msg[0]} {msg[1] or ''}" for msg in history]) if history else ''
+        full_prompt = f"{system_message}\n{full_history}\nUser: {message}\nAssistant:"
 
-    # Generate response
-    try:
-        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 or ""
-            response += token
-            yield response
-    except Exception as e:
-        yield f"Error during generation: {e}"
+        try:
+            response = pipe(
+                full_prompt, 
+                max_length=len(full_prompt) + max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                num_return_sequences=1
+            )[0]['generated_text']
+            
+            # Extract just the new assistant response
+            assistant_response = response[len(full_prompt):].strip()
+            return assistant_response
+        except Exception as e:
+            return f"Generation error: {e}"
+    
+    # For classification and sentiment models
+    elif model_name == "Text Classification":
+        try:
+            result = pipe(message)[0]
+            return f"Classification: {result['label']} (Confidence: {result['score']:.2f})"
+        except Exception as e:
+            return f"Classification error: {e}"
+    
+    elif model_name == "Sentiment Analysis":
+        try:
+            result = pipe(message)[0]
+            return f"Sentiment: {result['label']} (Confidence: {result['score']:.2f})"
+        except Exception as e:
+            return f"Sentiment analysis error: {e}"
 
 def create_chat_interface():
     """Create Gradio ChatInterface with model selection."""
     demo = gr.ChatInterface(
         respond,
         additional_inputs=[
-            gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-            gr.Dropdown(list(MODELS.keys()), value="Zephyr 7B Beta", label="Select Model"),
+            gr.Textbox(value="You are a helpful assistant.", label="System message"),
+            gr.Dropdown(
+                ["GPT-2 Original", "GPT-2 Medium", "DistilGPT-2", 
+                 "German GPT-2", "German Wechsel GPT-2", 
+                 "T5 Base", "T5 Large", 
+                 "Text Classification", "Sentiment Analysis"], 
+                value="GPT-2 Original", 
+                label="Select Model"
+            ),
             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(