Update app.py

app.py

@@ -1,38 +1,93 @@
+import os
+import gradio as gr
 from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
 
-app = FastAPI(title="Text Summarization API")
+os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers_cache'
+os.makedirs('/tmp/transformers_cache', exist_ok=True)
 
-summarizer = pipeline(
-    "summarization",
-    model="deepseek-ai/deepseek-llm-7b-base",
-    device=0 if torch.cuda.is_available() else -1
+app = FastAPI(title="DeepSeek LLM Interface")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+model_name = "deepseek-ai/deepseek-llm-7b-base"
+tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir='/tmp/transformers_cache')
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    cache_dir='/tmp/transformers_cache',
+    torch_dtype=torch.float16,
+    device_map="auto"
 )
 
-class SummarizationRequest(BaseModel):
-    text: str
-    max_length: int = 130
-    min_length: int = 20
+def generate_response(prompt, max_length=500, temperature=0.7):
+    """Generate response using the DeepSeek model"""
+    try:
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        outputs = model.generate(
+            **inputs,
+            max_length=max_length,
+            temperature=temperature,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response
+    except Exception as e:
+        print(f"Error in generate_response: {str(e)}")
+        return f"Error generating response: {str(e)}"
+
+class GenerationRequest(BaseModel):
+    prompt: str
+    max_length: int = 500
+    temperature: float = 0.7
 
-class SummarizationResponse(BaseModel):
-    summary: str
+class GenerationResponse(BaseModel):
+    response: str
 
-@app.post("/summarize", response_model=SummarizationResponse)
-async def summarize_text(request: SummarizationRequest):
+@app.post("/generate", response_model=GenerationResponse)
+async def generate_text(request: GenerationRequest):
     try:
-        summary = summarizer(
-            request.text,
+        response = generate_response(
+            request.prompt,
             max_length=request.max_length,
-            min_length=request.min_length,
-            do_sample=False
-        )[0]['summary_text']
-        
-        return SummarizationResponse(summary=summary)
+            temperature=request.temperature
+        )
+        return GenerationResponse(response=response)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.get("/health")
 async def health_check():
     return {"status": "healthy"}
+
+def gradio_generate(prompt, max_length, temperature):
+    return generate_response(prompt, int(max_length), float(temperature))
+
+interface = gr.Interface(
+    fn=gradio_generate,
+    inputs=[
+        gr.Textbox(label="Prompt", placeholder="Enter your prompt here..."),
+        gr.Slider(minimum=50, maximum=1000, value=500, step=50, label="Max Length"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
+    ],
+    outputs=gr.Textbox(label="Generated Response"),
+    title="DeepSeek LLM Interface",
+    description="Enter a prompt to generate text using DeepSeek LLM",
+    examples=[
+        ["Write a short story about a mysterious garden"],
+        ["Explain quantum computing in simple terms"],
+        ["Create a recipe for chocolate chip cookies"]
+    ]
+)
+
+app = gr.mount_gradio_app(app, interface, path="/")
+